Code review suggestions

Author: rumackaaron

changehc/delphi_changehc/download_ftp_files.py

@@ -10,39 +10,34 @@
 # third party
 import paramiko
 
-class AllowAnythingPolicy(paramiko.MissingHostKeyPolicy):
-    def missing_host_key(self, client, hostname, key):
-        return
-
 
 def print_callback(filename, bytes_so_far, bytes_total):
+    """Log file transfer progress"""
     rough_percent_transferred = int(100 * (bytes_so_far / bytes_total))
     if (rough_percent_transferred % 25) == 0:
         print(f'{filename} transfer: {rough_percent_transferred}%')
 
 
 def get_files_from_dir(sftp, out_path):
+    """Download files from sftp server that have been uploaded in last day
+    Args:
+        sftp: SFTP Session from Paramiko client
+        out_path: Path to local directory into which to download the files
+    """
+
     current_time = datetime.datetime.now()
-    seconds_in_day = 24 * 60 * 60
 
     # go through files in recieving dir
-    files_to_download = []
+    filepaths_to_download = {}
     for fileattr in sftp.listdir_attr():
         file_time = datetime.datetime.fromtimestamp(fileattr.st_mtime)
-        time_diff_to_current_time = current_time - file_time
-        if time_diff_to_current_time.total_seconds() <= seconds_in_day:
-            files_to_download.append(fileattr.filename)
-
-    filepaths_to_download = {}
-    for file in files_to_download:
-        full_path = path.join(out_path, file)
-        if path.exists(full_path):
-            print(f"{file} exists, skipping")
-        else:
-            filepaths_to_download[file] = full_path
+        filename = fileattr.filename
+        if current_time - file_time < datetime.timedelta(days=1) and \
+                not path.exists(filename):
+            filepaths_to_download[filename] = path.join(out_path, filename)
 
     # make sure we don't download more than 2 files per day
-    assert len(files_to_download) <= 2, "more files dropped than expected"
+    assert len(filepaths_to_download) <= 2, "more files dropped than expected"
 
     # download!
     for infile, outfile in filepaths_to_download.items():
@@ -51,21 +46,29 @@ def get_files_from_dir(sftp, out_path):
 
 
 def download(out_path, ftp_conn):
+    """Downloads files necessary to create CHC signal from ftp server.
+    Args:
+        out_path: Path to local directory into which to download the files
+        ftp_conn: Dict containing login credentials to ftp server
+    """
 
     # open client
-    client = paramiko.SSHClient()
-    client.set_missing_host_key_policy(AllowAnythingPolicy())
-
-    client.connect(ftp_conn["host"], username=ftp_conn["user"],
-                   password=ftp_conn["pass"][1:] + ftp_conn["pass"][0],
-                   port=ftp_conn["port"],
-                   allow_agent=False, look_for_keys=False)
-    sftp = client.open_sftp()
-
-    sftp.chdir('/dailycounts/All_Outpatients_By_County')
-    get_files_from_dir(sftp, out_path)
-
-    sftp.chdir('/dailycounts/Covid_Outpatients_By_County')
-    get_files_from_dir(sftp, out_path)
-
-    client.close()
+    try:
+        client = paramiko.SSHClient()
+        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+        client.connect(ftp_conn["host"], username=ftp_conn["user"],
+                       password=ftp_conn["pass"][1:] + ftp_conn["pass"][0],
+                       port=ftp_conn["port"],
+                       allow_agent=False, look_for_keys=False)
+        sftp = client.open_sftp()
+
+        sftp.chdir('/dailycounts/All_Outpatients_By_County')
+        get_files_from_dir(sftp, out_path)
+
+        sftp.chdir('/dailycounts/Covid_Outpatients_By_County')
+        get_files_from_dir(sftp, out_path)
+    
+    finally:
+        if client:
+            client.close()