fix(parquet): include parsings to parquet reading

Author: luabida

pysus/api/client.py

@@ -10,7 +10,9 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Literal
 
+import anyio
 import duckdb
+import pandas as pd
 from pysus import CACHEPATH
 from sqlalchemy import DateTime, Enum, Integer, String, create_engine
 from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, sessionmaker
@@ -235,8 +237,17 @@ async def download(
         file: BaseRemoteFile,
         token: str | None = None,
         callback: Callable | None = None,
+        timeout: float | None = None,
     ) -> BaseLocalFile:
-        """Download a remote file and return a local file handle."""
+        """Download a remote file and return a local file handle.
+
+        Parameters
+        ----------
+        timeout : float | None
+            Maximum seconds to wait for the download. ``None`` (default) means
+            no timeout – use this when the socket-level timeout on the
+            underlying client is sufficient.
+        """
 
         from pysus.api.extensions import ExtensionFactory
 
@@ -271,7 +282,11 @@ async def download(
                     f"No download logic for client: {client_name}",
                 )
 
-            await client._download_file(file, local_path, callback)
+            if timeout is not None:
+                with anyio.fail_after(timeout):
+                    await client._download_file(file, local_path, callback)
+            else:
+                await client._download_file(file, local_path, callback)
 
             await self._update_state(
                 local_path=local_path,
@@ -311,18 +326,22 @@ async def download_to_parquet(
         file: BaseRemoteFile,
         token: str | None = None,
         callback: Callable[[int, int], None] | None = None,
+        timeout: float | None = None,
+        add_dv: bool = True,
     ) -> Parquet:
         """Download a file and convert it to Parquet format."""
 
         local_file = await self.download(
             file=file,
             token=token,
             callback=callback,
+            timeout=timeout,
         )
 
         if hasattr(local_file, "to_parquet"):
             original_path = local_file.path
             parquet_file = await local_file.to_parquet(callback=callback)
+            parquet_file.add_dv = add_dv
 
             await self._update_state(
                 local_path=parquet_file.path,
@@ -346,7 +365,9 @@ async def download_to_parquet(
         )
 
     def get_local_hierarchy(self):
-        """Build a nested dict of cached files grouped by client and dataset."""
+        """
+        Build a nested dict of cached files grouped by client and dataset.
+        """
 
         with self.Session() as session:
             records = session.query(LocalFileState).all()
@@ -414,8 +435,20 @@ def read_parquet(
         paths: list[Path],
         sql: str | None = None,
         mode: Literal["union", "intersection", "strict"] = "union",
-    ) -> "DuckDBPyConnection":
-        """Read Parquet files with optional schema handling and SQL filter."""
+        add_dv: bool = True,
+    ) -> "DuckDBPyConnection | pd.DataFrame":
+        """Read Parquet files with optional schema handling and SQL filter.
+
+        Parameters
+        ----------
+        add_dv : bool
+            When True, automatically applies the IBGE verification digit to
+            municipality code columns. If there are matching columns, a
+            DataFrame is returned instead of a DuckDBPyConnection.
+        """
+
+        from pysus.api.utils import add_dv as _add_dv_fn
+        from pysus.api.utils import is_geocode_column
 
         if not paths:
             raise ValueError("No paths provided")
@@ -452,8 +485,7 @@ def get_columns(path: Path) -> set[tuple[str, str]]:
         else:
             paths_str = ", ".join(f"'{p}'" for p in paths)
             query = (
-                f"SELECT * FROM read_parquet([{paths_str}], "
-                "union_by_name=True)"
+                f"SELECT * FROM read_parquet([{paths_str}], union_by_name=True)"
             )
 
         if sql:
@@ -462,4 +494,29 @@ def get_columns(path: Path) -> set[tuple[str, str]]:
             else:
                 query = f"SELECT {sql} FROM ({query}) AS t"
 
+        base = duckdb.execute(query)
+
+        if not add_dv:
+            return base
+
+        geocode_cols = [
+            col[0] for col in base.description if is_geocode_column(col[0])
+        ]
+        if not geocode_cols:
+            return base
+
+        duckdb.create_function(
+            "__pysus_add_dv",
+            _add_dv_fn,
+            null_handling="special",
+        )
+        selects = [
+            (
+                f'__pysus_add_dv("{c[0]}") AS "{c[0]}"'
+                if c[0] in geocode_cols
+                else f'"{c[0]}"'
+            )
+            for c in base.description
+        ]
+        query = f"SELECT {', '.join(selects)} FROM ({query}) AS _t"
         return duckdb.execute(query)

pysus/api/extensions.py

@@ -188,6 +188,7 @@ class Parquet(BaseTabularFile):
     """Represents a Parquet file with optional date and integer type parsing."""
 
     type: FileType = Field("PARQUET")
+    add_dv: bool = True
 
     @property
     def schema(self) -> pa.Schema:
@@ -204,12 +205,26 @@ def rows(self) -> int:
         """Return the number of rows from the Parquet metadata."""
         return pq.read_metadata(self.path).num_rows
 
+    @staticmethod
+    def _apply_add_dv(df: pd.DataFrame) -> pd.DataFrame:
+        """Apply the IBGE verification digit to geocode columns in-place."""
+        from pysus.api.utils import add_dv, is_geocode_column
+
+        geocode_cols = [c for c in df.columns if is_geocode_column(c)]
+        for col in geocode_cols:
+            df[col] = df[col].astype(str).apply(add_dv)
+        return df
+
     async def load(self, parse: bool = True) -> pd.DataFrame:
         """Read the entire Parquet file into a DataFrame."""
 
         def _load():
             df = pd.read_parquet(self.path, engine="pyarrow")
-            return self.parse_dftypes(df) if parse else df
+            if parse:
+                df = self.parse_dftypes(df)
+                if self.add_dv:
+                    df = self._apply_add_dv(df)
+            return df
 
         return await to_thread.run_sync(_load)
 
@@ -226,6 +241,8 @@ async def stream(
             df = batch.to_pandas()
             if parse:
                 df = self.parse_dftypes(df)
+                if self.add_dv:
+                    df = self._apply_add_dv(df)
             yield df
             await asyncio.sleep(0)
 

pysus/api/ftp/client.py

@@ -43,6 +43,7 @@ class FTP(BaseRemoteClient):
     """Async FTP client for navigating and downloading DATASUS data."""
 
     host: str = "ftp.datasus.gov.br"
+    timeout: int = 60
 
     _ftp: FTPLib | None = PrivateAttr(default=None)
 
@@ -77,7 +78,7 @@ async def connect(self) -> None:
 
         def _connect():
             if self.ftp is None:
-                self._ftp = FTPLib(self.host)
+                self._ftp = FTPLib(self.host, timeout=self.timeout)
                 self.ftp.login()
 
         await to_thread.run_sync(_connect)

pysus/api/utils.py

@@ -0,0 +1,51 @@
+GEOCODE_PREFIXES = (
+    "ID_MUNICIP",
+    "ID_MN_RESI",
+    "ID_MUNI_RE",
+    "MUN_",
+    "COD_MUN_",
+    "CO_MUN_",
+    "ID_MUNI_AT",
+    "ID_MUNIC_",
+)
+
+
+def is_geocode_column(name: str) -> bool:
+    """Check if a column name corresponds to an IBGE municipality code."""
+    upper = name.upper()
+    return any(upper.startswith(p) for p in GEOCODE_PREFIXES)
+
+
+def add_dv(geocode: str) -> str:
+    if not geocode or not str(geocode).isdigit():
+        return geocode
+
+    miscalculated = {
+        "2201911": "2201919",
+        "2201986": "2201988",
+        "2202257": "2202251",
+        "2611531": "2611533",
+        "3117835": "3117836",
+        "3152139": "3152131",
+        "4305876": "4305871",
+        "5203963": "5203962",
+        "5203930": "5203939",
+    }
+
+    if len(str(geocode)) == 7:
+        return miscalculated.get(str(geocode), geocode)
+
+    if len(str(geocode)) == 6:
+        weight = [1, 2, 1, 2, 1, 2]
+        total = sum(
+            sum(divmod(int(d) * w, 10))
+            for d, w in zip(
+                str(geocode),
+                weight,
+            )
+        )
+        dv = 0 if total % 10 == 0 else 10 - (total % 10)
+        code = str(geocode) + str(dv)
+        return miscalculated.get(code, code)
+
+    return geocode

pysus/tests/api/ftp/test_client.py

@@ -48,7 +48,9 @@ async def test_connect_and_login(ftp_client):
         mock_instance = mock_ftplib.return_value
         await ftp_client.login()
 
-        mock_ftplib.assert_called_once_with(ftp_client.host)
+        mock_ftplib.assert_called_once_with(
+            ftp_client.host, timeout=ftp_client.timeout
+        )
         mock_instance.login.assert_called_once()