feat: adding constants

basedosdados · Oct 31, 2023 · 6f2cc3d · 6f2cc3d
1 parent 57acb45
commit 6f2cc3d
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 24 deletions.
diff --git a/pipelines/datasets/br_cgu_beneficios_cidadao/constants.py b/pipelines/datasets/br_cgu_beneficios_cidadao/constants.py
@@ -5,6 +5,8 @@
 
 from enum import Enum
 
+import numpy as np
+
 
 class constants(Enum):  # pylint: disable=c0103
     """
@@ -16,3 +18,25 @@ class constants(Enum):  # pylint: disable=c0103
     )
 
     MAIN_URL = "https://dadosabertos-download.cgu.gov.br/PortalDaTransparencia/saida/novo-bolsa-familia/"
+    DTYPES = {
+        "MÊS COMPETÊNCIA": str,
+        "MÊS REFERÊNCIA": str,
+        "UF": str,
+        "CÓDIGO MUNICÍPIO SIAFI": str,
+        "NOME MUNICÍPIO": str,
+        "CPF FAVORECIDO": str,
+        "NIS FAVORECIDO": str,
+        "NOME FAVORECIDO": str,
+        "VALOR PARCELA": np.float64,
+    }
+    RENAMER = {
+        "MÊS COMPETÊNCIA": "mes_competencia",
+        "MÊS REFERÊNCIA": "mes_referencia",
+        "UF": "sigla_uf",
+        "CÓDIGO MUNICÍPIO SIAFI": "id_municipio_siafi",
+        "NOME MUNICÍPIO": "municipio",
+        "CPF FAVORECIDO": "cpf",
+        "NIS FAVORECIDO": "nis",
+        "NOME FAVORECIDO": "nome",
+        "VALOR PARCELA": "valor",
+    }
diff --git a/pipelines/datasets/br_cgu_beneficios_cidadao/utils.py b/pipelines/datasets/br_cgu_beneficios_cidadao/utils.py
@@ -2,7 +2,6 @@
 import os
 import zipfile
 
-import dask.dataframe as dd
 import numpy as np
 import pandas as pd
 import requests
@@ -164,39 +163,19 @@ def parquet_partition(path):
             #    encoding="latin-1",
             #    dtype="string",
             # )
-            dtypes = {
-                "MÊS COMPETÊNCIA": str,
-                "MÊS REFERÊNCIA": str,
-                "UF": str,
-                "CÓDIGO MUNICÍPIO SIAFI": str,
-                "NOME MUNICÍPIO": str,
-                "CPF FAVORECIDO": str,
-                "NIS FAVORECIDO": str,
-                "NOME FAVORECIDO": str,
-                "VALOR PARCELA": np.float64,
-            }
+
             df = None
             with pd.read_csv(
                 f"{path}{nome_arquivo}",
                 sep=";",
                 encoding="latin-1",
-                dtype=dtypes,
+                dtype=constants.DTYPES.value,
                 chunksize=1000000,
                 decimal=",",
             ) as reader:
                 for chunk in tqdm(reader):
                     chunk.rename(
-                        columns={
-                            "MÊS COMPETÊNCIA": "mes_competencia",
-                            "MÊS REFERÊNCIA": "mes_referencia",
-                            "UF": "sigla_uf",
-                            "CÓDIGO MUNICÍPIO SIAFI": "id_municipio_siafi",
-                            "NOME MUNICÍPIO": "municipio",
-                            "CPF FAVORECIDO": "cpf",
-                            "NIS FAVORECIDO": "nis",
-                            "NOME FAVORECIDO": "nome",
-                            "VALOR PARCELA": "valor",
-                        },
+                        columns=constants.RENAMER.value,
                         inplace=True,
                     )
                     if df is None: