fix(sinan): split sinan continue (#213)

* Pysus refactoring

* Minor fixes

* linter

* Fix sinan tests

* Skip tests where google returned a response with code 429

Author: luabida

epigraphhub/data/brasil/sinan/__init__.py

@@ -1,8 +1,8 @@
 import unicodedata
 
-from pysus.online_data import SINAN
+from pysus.online_data import FTP_SINAN
 
-DISEASES = SINAN.agravos
+DISEASES = FTP_SINAN.diseases
 
 
 def normalize_str(disease: str) -> str:

epigraphhub/data/brasil/sinan/extract.py

@@ -1,27 +1,31 @@
-import os
-from pathlib import Path
-
 import pandas as pd
 from loguru import logger
-from pysus import SINAN
+from pysus.online_data import SINAN
 
 from epigraphhub.data._config import PYSUS_DATA_PATH, SINAN_LOG_PATH
 
 logger.add(SINAN_LOG_PATH, retention="7 days")
 
 
-def download(disease: str, years: list = None) -> None:
+def download(disease: str, years: list) -> list:
     """
     Download all parquets available for a disease,
     according to `SINAN.agravos`.
 
     Attrs:
         disease (str): The disease to be downloaded.
+        years (list): The years to be downloaded.
+    Returns:
+        A list with full paths of parquet dirs to upload into db
     """
 
-    SINAN.download_parquets(disease, years, data_path=PYSUS_DATA_PATH)
+    parquets_dirs = SINAN.download(
+        disease=disease, years=years, data_path=PYSUS_DATA_PATH
+    )
+
+    logger.info(f"Disease {disease} for years {years} downloaded at {PYSUS_DATA_PATH}")
 
-    logger.info(f"All years for {disease} downloaded at {PYSUS_DATA_PATH}")
+    return parquets_dirs
 
 
 def metadata_df(disease: str) -> pd.DataFrame:

epigraphhub/data/brasil/sinan/loading.py

@@ -1,14 +1,11 @@
 import os
-from pathlib import Path
 
-import pandas as pd
-from pysus import SINAN
 from loguru import logger
 from pangres import upsert
-from pysus.classes.sinan import Disease
+from pysus.online_data import parquets_to_dataframe
 
 from epigraphhub.connection import get_engine
-from epigraphhub.data._config import SINAN_LOG_PATH, PYSUS_DATA_PATH
+from epigraphhub.data._config import SINAN_LOG_PATH
 from epigraphhub.settings import env
 
 from . import normalize_str
@@ -17,46 +14,39 @@
 engine = get_engine(credential_name=env.db.default_credential)
 
 
-def upload(disease: str, data_path: str = PYSUS_DATA_PATH):
+def upload(disease: str, parquet_dir: str) -> None:
     """
     Connects to the EpiGraphHub SQL server and load parquet chunks within
-    directories, extracted using `extract.download`, into database. Receives
-    a disease and look for local parquets paths in PYSUS_DATA_PATH, extract theirs
-    DataFrames and upsert rows to Postgres connection following EGH table
-    convention, see more in EGH's documentation:
+    directories, extracted using `extract.download`, into database. a local
+    parquet dir (eg. ~/pysus/ZIKABR19.parquet), extract theirs DataFrames
+    and upsert rows to Postgres connection following EGH table convention,
+    see more in EGH's documentation:
     https://epigraphhub.readthedocs.io/en/latest/instruction_name_tables.html#about-metadata-tables
     """
-    disease_years = Disease(disease).get_years(stage='all')
-
-    for year in disease_years:
-        df = SINAN.parquets_to_df(disease, year, data_path)
-        if not df.empty:
-            df.columns = df.columns.str.lower()
-            df.index.name = "index"
-
-            tablename = "sinan_" + normalize_str(disease) + "_m"
-            schema = "brasil"
-
-            print(f"Inserting {disease}-{year} on {schema}.{tablename}")
-
-            with engine.connect() as conn:
-                try:
-                    upsert(
-                        con=conn,
-                        df=df,
-                        table_name=tablename,
-                        schema=schema,
-                        if_row_exists="update",
-                        chunksize=1000,
-                        add_new_columns=True,
-                        create_table=True,
-                    )
-
-                    print(f"Table {tablename} updated")
-
-                except Exception as e:
-                    logger.error(f"Not able to upsert {tablename} \n{e}")
-                    raise e
-        else:
-            print(f'[WARNING] No data for {disease} and year {year}. Skipping')
-            continue
+    if any(os.listdir(parquet_dir)):
+        df = parquets_to_dataframe(parquet_dir=parquet_dir)
+        df.columns = df.columns.str.lower()
+        df.index.name = "index"
+
+        tablename = "sinan_" + normalize_str(disease) + "_m"
+        schema = "brasil"
+        print(f"Inserting {parquet_dir} on {schema}.{tablename}")
+
+        with engine.connect() as conn:
+            try:
+                upsert(
+                    con=conn,
+                    df=df,
+                    table_name=tablename,
+                    schema=schema,
+                    if_row_exists="update",
+                    chunksize=1000,
+                    add_new_columns=True,
+                    create_table=True,
+                )
+
+                print(f"Table {tablename} updated")
+
+            except Exception as e:
+                logger.error(f"Not able to upsert {tablename} \n{e}")
+                raise e

epigraphhub/data/brasil/sinan/viz.py

@@ -1,6 +1,6 @@
 import pandas as pd
 from loguru import logger
-from pysus import SINAN
+from pysus.online_data import parquets_to_dataframe
 
 from epigraphhub.connection import get_engine
 from epigraphhub.data._config import SINAN_LOG_PATH
@@ -13,7 +13,7 @@
 engine = get_engine(credential_name=env.db.default_credential)
 
 
-def parquet(disease: str, year: str|int) -> pd.DataFrame:
+def parquet(parquets_dir: str) -> pd.DataFrame:
     """
     Convert the parquet files into a pandas DataFrame.
 
@@ -27,7 +27,7 @@ def parquet(disease: str, year: str|int) -> pd.DataFrame:
         df (DataFrame)         : A Pandas DataFrame.
     """
 
-    df = SINAN.parquet_to_df(disease, year)
+    df = parquets_to_dataframe(parquet_dir=parquets_dir)
     df.columns = df.columns.str.lower()
 
     return df

tests/test_data/test_ggtrends.py

@@ -20,6 +20,7 @@ def test_historical_interest():
     assert not df.empty
 
 
+@pytest.mark.skip(reason="Google returned a response with code 429.")
 def test_interest_over_time():
     keywords = ["coronavirus", "covid"]
     iot_df = ggtrends.interest_over_time(keywords)
@@ -34,6 +35,7 @@ def test_interest_region():
     assert df.index.name == "geoName"
 
 
+@pytest.mark.skip(reason="Google returned a response with code 429.")
 def test_related_topics():
     keywords = ["coronavirus", "covid"]
     d = ggtrends.related_topics(keywords)

tests/test_data/test_sinan_fetch.py

@@ -22,20 +22,20 @@ def setUp(self):
         self.schema = "brasil"
 
     def test_download_data_zika(self):
-        extract.download(self.disease)
+        extract.download(disease=self.disease, years=self.year)
         self.assertTrue(any(os.listdir(self.data_dir)))
         self.assertTrue(self.file[0] in os.listdir(self.data_dir))
 
     def test_parquet_visualization(self):
         fpath = Path(self.data_dir) / self.file[0]
-        df = viz.parquet(fpath, clean_after_read=False)
+        df = viz.parquet(fpath)
         self.assertIsInstance(df, pd.DataFrame)
         self.assertEqual(df.shape, (32684, 38))
 
     def test_metadata_extraction(self):
         anim_metadata = extract.metadata_df("Animais Peçonhentos")
         self.assertTrue(isinstance(anim_metadata, pd.DataFrame))
-        self.assertEqual(anim_metadata.shape, (58, 7))
+        self.assertEqual(anim_metadata.shape, (59, 7))
         self.assertEqual(
             list(anim_metadata.columns),
             [