Fixed many bugs.

testing passing for sources POPTCU and projpop

pysus/ftp/databases/ibge.py → pysus/ftp/databases/ibge_datasus.py

@@ -3,14 +3,15 @@
 from pysus.ftp import Database, Directory, File
 from pysus.ftp.utils import zfill_year, to_list
 
+
 class IBGEDATASUS(Database):
     name = "IBGE-DataSUS"
     paths = (
         Directory("/dissemin/publicos/IBGE/POP"),
         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
-

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
-
-

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()

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()