Implement download methods on File class

pysus/ftp/__init__.py

@@ -2,7 +2,10 @@
 from ftplib import FTP
 import os
 import pathlib
-from typing import List, Union, Optional
+from aioftp import Client
+from typing import List, Optional, Union
+
+from pysus.online_data import CACHEPATH
 
 
 class File:
@@ -50,9 +53,55 @@ def __eq__(self, other):
             return self.path == other.path
         return False
 
-    def download(self):  # -> task.run
-        """TODO"""
-        ...
+    def download(self, local_dir: str = CACHEPATH):
+        dir = pathlib.Path(local_dir)
+        dir.mkdir(exist_ok=True, parents=True)
+        filepath = dir / self.basename
+
+        if filepath.exists():
+            pass
+
+        ftp = ftp = FTP("ftp.datasus.gov.br")
+        ftp.login()
+        ftp.retrbinary(
+            f"RETR {self.path}",
+            open(f"{filepath}", "wb").write,
+        )
+        ftp.close()
+        return str(filepath)
+
+    async def async_download(self, local_dir: str = CACHEPATH):
+        # aioftp.Client.parse_list_line_custom
+        def line_file_parser(file_line):
+            line = file_line.decode("utf-8")
+            info = {}
+            if "<DIR>" in line:
+                date, time, _, *name = str(line).strip().split()
+                info["size"] = 0
+                info["type"] = "dir"
+                name = " ".join(name)
+            else:
+                date, time, size, name = str(line).strip().split()
+                info["size"] = size
+                info["type"] = "file"
+
+            modify = datetime.strptime(
+                " ".join([date, time]), "%m-%d-%y %I:%M%p"
+            )
+            info["modify"] = modify.strftime("%m/%d/%Y %I:%M%p")
+
+            return pathlib.PurePosixPath(name), info
+
+        output = (
+            local_dir + str(self.basename)
+            if local_dir.endswith("/")
+            else local_dir + "/" + str(self.basename)
+        )
+        async with Client.context(
+            host="ftp.datasus.gov.br", parse_list_line_custom=line_file_parser
+        ) as client:
+            await client.login()
+            await client.download(self.path, output, write_into=True)
 
 
 class Directory:
@@ -90,25 +139,6 @@ def __eq__(self, other):
         return False
 
 
-# aioftp.Client.parse_list_line_custom
-def line_file_parser(file_line):
-    info = {}
-    if "<DIR>" in file_line:
-        date, time, _, *name = str(file_line).strip().split()
-        info["size"] = 0
-        info["type"] = "dir"
-        name = " ".join(name)
-    else:
-        date, time, size, name = str(file_line).strip().split()
-        info["size"] = size
-        info["type"] = "file"
-
-    modify = datetime.strptime(" ".join([date, time]), "%m-%d-%y %I:%M%p")
-    info["modify"] = modify.strftime("%m/%d/%Y %I:%M%p")
-
-    return pathlib.PurePosixPath(name), info
-
-
 def list_path(path: str) -> List[Union[Directory, File]]:
     """
     This method is responsible for listing all the database's

pysus/ftp/databases.py

@@ -116,16 +116,19 @@ class SINAN(Database):
     )
 
     def describe(self, file: File) -> dict:
-        dis_code, year = self.format(file)
+        if file.extension.upper() == ".DBC":
+            dis_code, year = self.format(file)
 
-        description = dict(
-            name=str(file.basename),
-            disease=self.diseases[dis_code],
-            year=zfill_year(year),
-            size=humanize.naturalsize(file.info["size"]),
-            last_update=file.info["modify"].strftime("%m-%d-%Y %I:%M%p"),
-        )
-        return description
+            description = dict(
+                name=str(file.basename),
+                disease=self.diseases[dis_code],
+                year=zfill_year(year),
+                size=humanize.naturalsize(file.info["size"]),
+                last_update=file.info["modify"].strftime("%m-%d-%Y %I:%M%p"),
+            )
+            return description
+        else:
+            return {}
 
     def format(self, file: File) -> tuple:
         year = file.name[-2:]
@@ -270,11 +273,8 @@ class SINASC(Database):
         DNR="Dados dos Nascidos Vivos por UF de residência",
     )
 
-    def describe(self, files: Union[File, list[File]]) -> dict:
-        files = to_list(files)
-        description = dict()
-
-        for file in files:
+    def describe(self, file: File) -> dict:
+        if file.extension.upper() == ".DBC":
             uf, year = self.format(file)
 
             if uf == "EX":  # DNEX2021.dbc
@@ -291,6 +291,8 @@ def describe(self, files: Union[File, list[File]]) -> dict:
             )
 
             return description
+        else:
+            return {}
 
     def format(self, file: File) -> tuple:
         if file.name == "DNEX2021":