Started the implementation of fetching the population estimates from …

.gitignore

@@ -1,4 +1,4 @@
-.idea
+.idea/
 PySUS.egg-info/
 build/
 dist/

docs/source/conf.py

@@ -162,7 +162,7 @@
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,
 # so a file named "default.css" will overwrite the builtin "default.css".
-html_static_path = ["_static"]
+# html_static_path = ["_static"]
 
 # Add any extra paths that contain custom files (such as robots.txt or
 # .htaccess) here, relative to this directory. These files are copied

pyproject.toml

@@ -39,6 +39,7 @@ flake8 = "^5.0.4"
 isort = "^5.10.1"
 pre-commit = "^2.20.0"
 pytest-timeout = "^2.1.0"
+nbsphinx = "^0.9.3"
 
 [tool.poetry.group.docs.dependencies]
 sphinx = "^5.1.1"

pysus/ftp/databases/ibge_datasus.py

@@ -0,0 +1,67 @@
+from typing import Optional, List, Union
+
+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") # this has a different file name pattern
+    )
+    metadata = {
+        "long_name": "Populaçao Residente, Censos, Contagens "
+                     "Populacionais e Projeçoes Intercensitarias",
+        "source": "ftp://ftp.datasus.gov.br/dissemin/publicos/IBGE",
+        "description": (
+            "São aqui apresentados informações sobre a população residente, "
+            "estratificadas por município, faixas etárias e sexo, obtidas a "
+            "partir dos Censos Demográficos, Contagens Populacionais "
+            "e Projeções Intercensitárias."
+        ),
+    }
+
+    def describe(self, file: File) -> dict:
+        if file.extension.upper() in [".ZIP"]:
+            year = file.name.split('.')[0][-2:]
+            description = {
+                "name": str(file.basename),
+                "year": zfill_year(year),
+                "size": file.info["size"],
+                "last_update": file.info["modify"]
+            }
+            return description
+        elif file.extension.upper() == ".DBF":
+            year = file.name[-2:]
+            description = {
+                "name": str(file.basename),
+                "year": zfill_year(year),
+                "size": file.info["size"],
+                "last_update": file.info["modify"]
+            }
+            return description
+        return {}
+
+    def format(self, file: File) -> str:
+        return file.name[-2:]
+
+    def get_files(
+            self,
+            year: Optional[Union[str, int, list]] = None,
+    ) -> List[File]:
+        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: zfill_year(self.format(f)) in years, files))
+
+        return files

pysus/online_data/IBGE.py

@@ -6,6 +6,8 @@
 import requests
 import pandas as pd
 
+from pysus.ftp.databases.ibge_datasus import IBGEDATASUS
+
 # requests.packages.urllib3.util.ssl_.DEFAULT_CIPHERS = 'ALL:@SECLEVEL=1'
 
 from urllib.error import HTTPError
@@ -14,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
@@ -230,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
@@ -288,3 +290,15 @@ def get_legacy_session():
     session = requests.session()
     session.mount('https://', CustomHttpAdapter(ctx))
     return session
+
+
+def get_population(year, source='POPTCU'):
+    """
+    Get population data from IBGE as shared by DATASUS
+    :param year: year of the data
+    :param source: 'POPTCU'|'POP'|'censo'|'projpop'
+    :return: DataFrame with population data
+    """
+    ibgedatasus = IBGEDATASUS().load()
+    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()