diff --git a/pysus/ftp/__init__.py b/pysus/ftp/__init__.py index d8b1389f..cf787c14 100644 --- a/pysus/ftp/__init__.py +++ b/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 "" 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 "" 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 diff --git a/pysus/ftp/async_utils.py b/pysus/ftp/async_utils.py deleted file mode 100644 index 9d6ce02c..00000000 --- a/pysus/ftp/async_utils.py +++ /dev/null @@ -1,18 +0,0 @@ -from aioftp import Client - -from . import File, line_file_parser -from pysus.online_data import CACHEPATH - - -async def download(file: File, local_dir: str = CACHEPATH): - output = ( - local_dir+file.basename - if local_dir.endswith("/") - else local_dir+"/"+file.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(file.path, output, write_into=True) diff --git a/pysus/ftp/databases.py b/pysus/ftp/databases.py index f06462e2..ae526008 100644 --- a/pysus/ftp/databases.py +++ b/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":