diff --git a/pysus/ftp/databases/ibge.py b/pysus/ftp/databases/ibge_datasus.py similarity index 75% rename from pysus/ftp/databases/ibge.py rename to pysus/ftp/databases/ibge_datasus.py index 4e1f34e..dd0aa7f 100644 --- a/pysus/ftp/databases/ibge.py +++ b/pysus/ftp/databases/ibge_datasus.py @@ -3,6 +3,7 @@ from pysus.ftp import Database, Directory, File from pysus.ftp.utils import zfill_year, to_list + class IBGEDATASUS(Database): name = "IBGE-DataSUS" paths = ( @@ -10,7 +11,7 @@ class IBGEDATASUS(Database): Directory("/dissemin/publicos/IBGE/censo"), Directory("/dissemin/publicos/IBGE/POPTCU"), Directory("/dissemin/publicos/IBGE/projpop"), - Directory("/dissemin/publicos/IBGE/Auxiliar") + # Directory("/dissemin/publicos/IBGE/Auxiliar") # this has a different file name pattern ) metadata = { "long_name": "Populaçao Residente, Censos, Contagens " @@ -24,10 +25,9 @@ class IBGEDATASUS(Database): ), } - def describe(self, file: File) -> dict: - if file.extension.upper() == ".ZIP": - year = file.name[-2:] + if file.extension.upper() in [".ZIP"]: + year = file.name.split('.')[0][-2:] description = { "name": str(file.basename), "year": zfill_year(year), @@ -47,21 +47,21 @@ def describe(self, file: File) -> dict: return {} def format(self, file: File) -> str: - return file.name[-4:] + return file.name[-2:] def get_files( - self, - year: Optional[Union[str, int, list]] = None, + self, + year: Optional[Union[str, int, list]] = None, ) -> List[File]: - files = list(filter( - lambda f: f.extension.upper() in [".ZIP"], self.files - )) + files = [f for f in self.files if f.extension.upper() in [".ZIP", ".DBF"] and self.describe(f)["year"] == year] + # files = list(filter( + # lambda f: f.extension.upper() in [".ZIP"], self.files + # )) if year or str(year) in ["0", "00"]: years = ( [zfill_year(str(y)[-4:]) for y in to_list(year)] ) - files = list(filter(lambda f: self.format(f) in years, files)) + files = list(filter(lambda f: zfill_year(self.format(f)) in years, files)) return files - diff --git a/pysus/online_data/IBGE.py b/pysus/online_data/IBGE.py index 8bd240b..b209c70 100644 --- a/pysus/online_data/IBGE.py +++ b/pysus/online_data/IBGE.py @@ -6,7 +6,7 @@ import requests import pandas as pd -from pysus.ftp.database.ibge import IBGEDATASUS +from pysus.ftp.databases.ibge_datasus import IBGEDATASUS # requests.packages.urllib3.util.ssl_.DEFAULT_CIPHERS = 'ALL:@SECLEVEL=1' @@ -16,16 +16,16 @@ def get_sidra_table( - table_id, - territorial_level, - geocode='all', - period=None, - variables=None, - classification=None, - categories=None, - format=None, - decimals=None, - headers=None, + table_id, + territorial_level, + geocode='all', + period=None, + variables=None, + classification=None, + categories=None, + format=None, + decimals=None, + headers=None, ): """ Wrapper for the SIDRA API. More information here: http://apisidra.ibge.gov.br/home/ajuda @@ -232,11 +232,11 @@ class FetchData: """ def __init__( - self, agregado: int, periodos: str, variavel: str = 'allxp', **kwargs + self, agregado: int, periodos: str, variavel: str = 'allxp', **kwargs ): self.url = ( - APIBASE - + f'agregados/{agregado}/periodos/{periodos}/variaveis/{variavel}?' + APIBASE + + f'agregados/{agregado}/periodos/{periodos}/variaveis/{variavel}?' ) self.url += '&'.join([f'{k}={v}' for k, v in kwargs.items()]) self.JSON = None @@ -291,6 +291,7 @@ def get_legacy_session(): session.mount('https://', CustomHttpAdapter(ctx)) return session + def get_population(year, source='POPTCU'): """ Get population data from IBGE as shared by DATASUS @@ -299,7 +300,5 @@ def get_population(year, source='POPTCU'): :return: DataFrame with population data """ ibgedatasus = IBGEDATASUS().load() - files = [f for f in ibgedatasus.get_files(year=year) if f.path.endswith(source)] + files = [f for f in ibgedatasus.get_files(year=year) if f.path.split('/')[-2] == source] return files - - diff --git a/pysus/tests/test_ftp/test_databases/test_IBGEDATASUS.py b/pysus/tests/test_ftp/test_databases/test_IBGEDATASUS.py new file mode 100644 index 0000000..5e56b7d --- /dev/null +++ b/pysus/tests/test_ftp/test_databases/test_IBGEDATASUS.py @@ -0,0 +1,83 @@ +import unittest +from unittest.mock import MagicMock, patch +from pysus.ftp.databases import ibge_datasus + +class IBGEDATASUSTests(unittest.TestCase): + + @patch('pysus.ftp.databases.ibge_datasus.File') + def test_describe_zip_file(self, mock_file): + mock_file.extension.upper.return_value = ".ZIP" + mock_file.name = "POPBR12.zip" + mock_file.info = {"size": 100, "modify": "2022-01-01"} + + ibge = ibge_datasus.IBGEDATASUS() + result = ibge.describe(mock_file) + + self.assertEqual(result, { + "name": "POPBR12", + "year": "2012", + "size": 100, + "last_update": "2022-01-01" + }) + + # @patch('pysus.ftp.databases.ibge_datasus.File') + # def describe_dbf_file(self, mock_file): + # mock_file.extension.upper.return_value = ".DBF" + # mock_file.name = "file20.dbf" + # mock_file.info = {"size": 100, "modify": "2022-01-01"} + # + # ibge = ibge_datasus.IBGEDATASUS() + # result = ibge.describe(mock_file) + # + # self.assertEqual(result, { + # "name": "file20", + # "year": "2020", + # "size": 100, + # "last_update": "2022-01-01" + # }) + + # @patch('pysus.ftp.databases.ibge_datasus.File') + # def describe_other_file(self, mock_file): + # mock_file.extension.upper.return_value = ".TXT" + # mock_file.name = "file20.txt" + # + # ibge = ibge_datasus.IBGEDATASUS() + # result = ibge.describe(mock_file) + # + # self.assertEqual(result, {}) + + @patch('pysus.ftp.databases.ibge_datasus.File') + def format_file(self, mock_file): + mock_file.name = "file20.zip" + + ibge = ibge_datasus.IBGEDATASUS() + result = ibge.format(mock_file) + + self.assertEqual(result, "20.zip") + + @patch('pysus.ftp.databases.ibge_datasus.File') + @patch('pysus.ftp.databases.ibge_datasus.to_list') + def test_get_files_with_year(self, mock_to_list, mock_file): + mock_file.extension.upper.return_value = ".ZIP" + mock_file.name = "POPBR12.zip" + mock_to_list.return_value = ["2012"] + + ibge = ibge_datasus.IBGEDATASUS() + ibge.__content__ = {"POPBR12.zip": mock_file} + result = ibge.get_files(year="2012") + + self.assertEqual(result, [mock_file]) + + @patch('pysus.ftp.databases.ibge_datasus.File') + def get_files_without_year(self, mock_file): + mock_file.extension.upper.return_value = ".ZIP" + mock_file.name = "file20.zip" + + ibge = ibge_datasus.IBGEDATASUS() + ibge.files = [mock_file] + result = ibge.get_files() + + self.assertEqual(result, [mock_file]) + +if __name__ == '__main__': + unittest.main() \ No newline at end of file diff --git a/pysus/tests/test_ibge.py b/pysus/tests/test_ibge.py index 3948dd7..778fe4d 100644 --- a/pysus/tests/test_ibge.py +++ b/pysus/tests/test_ibge.py @@ -44,6 +44,15 @@ def test_FetchData(self): self.assertIsInstance(ds, IBGE.FetchData) self.assertGreater(len(ds.JSON), 0) + @pytest.mark.timeout(120) + def test_get_population(self): + l = IBGE.get_population(2021) + self.assertEqual(l[0].name, 'POPTBR21') + self.assertGreater(len(l), 0) + l = IBGE.get_population(2012, source='projpop') + self.assertEqual(l[0].name, 'projbr12') + self.assertGreater(len(l), 0) + if __name__ == '__main__': unittest.main()