feat: include a python API to dados.gov.br API

.gitignore

@@ -179,6 +179,7 @@ dmypy.json
 
 # pytype static type analyzer
 .pytype/
+.pylintrc
 
 # Cython debug symbols
 cython_debug/

condarecipe/pysus/meta.yaml

@@ -32,7 +32,6 @@ requirements:
     - pyarrow
     - python
     - requests
-    - elasticsearch
 
 test:
   imports:

pyproject.toml

@@ -2,10 +2,17 @@
 name = "pysus"
 version = "1.0.1"  # changed by semantic-release
 description = "Tools for dealing with Brazil's Public health data"
-authors = ["Flavio Codeco Coelho <fccoelho@gmail.com>"]
+authors = ["Flavio Codeco Coelho <fccoelho@gmail.com>", "Luã Bida Vacaro <luabidaa@gmail.com>"]
 license = "GPL"
 
-packages = [{include='pysus'}]
+packages = [{ include = "pysus"}]
+
+exclude = [
+  "pysus/tests",
+  "pysus/tests/**",
+  "pysus/management",
+  "pysus/management/**"
+]
 
 [tool.poetry.dependencies]
 python = ">=3.10,<3.14"
@@ -23,12 +30,15 @@ Unidecode = "^1.3.6"
 dateparser = "^1.1.8"
 pandas = "^2.2.2"
 urwid = "^2.1.2"
-elasticsearch = { version = "7.16.2", extras=["preprocessing"] }
 # FTP
 bigtree = "^0.12.2"
 aioftp = "^0.21.4"
 humanize = "^4.8.0"
-typing-extensions = "^4.9.0"
+typing-extensions = ">=4.10.0"
+pydantic = "^2.12.5"
+duckdb = "^1.4.4"
+duckdb-engine = "^0.17.0"
+sqlalchemy = "^2.0.48"
 
 [tool.poetry.group.dev.dependencies]
 pytest = ">=6.1.0"
@@ -39,6 +49,7 @@ pre-commit = "^2.20.0"
 pytest-timeout = "^2.1.0"
 nbsphinx = "^0.9.3"
 pytest-retry = "1.7.0"
+datamodel-code-generator = {extras = ["http"], version = "^0.54.0"}
 
 [tool.poetry.group.docs.dependencies]
 sphinx = "^5.1.1"

pysus/__init__.py

@@ -1,24 +1,25 @@
-# type: ignore[attr-defined]
 """PySUS Python package"""
 
+import os
+import pathlib
+from typing import Final
 from importlib import metadata as importlib_metadata
 
-from pysus.ftp.databases import *  # noqa
-from pysus.ftp.databases import AVAILABLE_DATABASES
+
+CACHEPATH: Final[str] = os.getenv(
+    "PYSUS_CACHEPATH",
+    os.path.join(str(pathlib.Path.home()), "pysus"),
+)
+
+from pysus.api.ftp.databases import *  # noqa
 
 
 def get_version() -> str:
     try:
         return importlib_metadata.version(__name__)
     except importlib_metadata.PackageNotFoundError:  # pragma: no cover
-        return "1.0.1"  # changed by semantic-release
+        return "1.0.1"
 
 
 version: str = get_version()
 __version__: str = version
-
-__all__ = [
-    "AVAILABLE_DATABASES",
-    "version",
-    "__version__",
-]

pysus/api/dadosgov/client.py

@@ -0,0 +1,53 @@
+import requests
+from typing import List, Optional
+from pydantic import TypeAdapter
+from pysus.api.dadosgov.models import (
+    Dataset,
+    DatasetSummary,
+)
+from pysus import __version__
+
+
+class DadosGov:
+    def __init__(self, token: str):
+        self.base_url = "https://dados.gov.br/dados/api"
+        self.session = requests.Session()
+        self.session.headers.update(
+            {
+                "Accept": "application/json",
+                "User-Agent": f"PySUS/{__version__}",
+                "chave-api-dados-abertos": token,
+            }
+        )
+
+    def _get(self, endpoint: str, params: Optional[dict] = None):
+        url = f"{self.base_url}/{endpoint.lstrip('/')}"
+        response = self.session.get(url, params=params)
+        response.raise_for_status()
+        return response.json()
+
+    def list_datasets(
+        self,
+        pagina: int = 1,
+        nome_conjunto: Optional[str] = None,
+        dados_abertos: Optional[bool] = None,
+        is_privado: bool = False,
+        id_organizacao: Optional[str] = None,
+    ) -> List[DatasetSummary]:
+        params = {
+            "pagina": pagina,
+            "nomeConjuntoDados": nome_conjunto,
+            "dadosAbertos": dados_abertos,
+            "isPrivado": is_privado,
+            "idOrganizacao": id_organizacao,
+        }
+
+        params = {k: v for k, v in params.items() if v is not None}
+
+        data = self._get("/publico/conjuntos-dados", params=params)
+        adapter = TypeAdapter(List[DatasetSummary])
+        return adapter.validate_python(data)
+
+    def get_dataset(self, id: str) -> Dataset:
+        data = self._get(f"/publico/conjuntos-dados/{id}")
+        return Dataset.model_validate(data)

pysus/api/dadosgov/databases.py

@@ -0,0 +1,90 @@
+__all__ = [
+    "CNES",
+    "PNI",
+    "SIA",
+    "SINAN",
+]
+
+from typing import List, Optional, Union
+
+from .models import Dataset, Resource
+from pysus.utils import UFs, parse_UFs, to_list, zfill_year, MONTHS
+from pysus.api.models import FileDescription
+
+
+class CNES(Dataset):
+    name = "CNES"
+    ids = (
+        "40a0d093-b12f-44a4-bdc7-bae8eb54dd04",
+        "9455b341-b06e-408e-8e10-54b32b3d74ec",
+    )
+
+    def describe(self, file: Resource) -> Optional[FileDescription]: ...
+
+    def format(self, file: Resource) -> tuple: ...
+
+    def get_files(
+        self,
+        year: Optional[Union[list, str, int]] = None,
+        month: Optional[Union[list, str, int]] = None,
+    ) -> List[Resource]: ...
+
+
+class PNI(Dataset):
+    name = "PNI"
+    ids = (
+        "2989d396-cb09-47e7-a3b8-a4b951ca0200",
+        "543aa08a-46c4-44e8-802e-198daa30753d",
+        "04292d08-ee4f-463a-b7b5-76cfb76775b3",
+        "7ed6eecc-c254-475c-92c5-daba5727596b",
+        "783b7456-6a6c-4025-a8bd-8e9caa0fb962",
+        "c6c3c6f3-2026-48a2-84ac-d8039714a0ba",
+        "9a25b796-80e3-444a-a4e7-405f5596d8ab",
+    )
+
+    def describe(self, file: Resource) -> Optional[FileDescription]: ...
+
+    def format(self, file: Resource) -> tuple: ...
+
+    def get_files(
+        self,
+        year: Optional[Union[list, str, int]] = None,
+        month: Optional[Union[list, str, int]] = None,
+    ) -> List[Resource]: ...
+
+
+class SIA(Dataset):
+    name = "SIA"
+    ids = ("9a335cb7-2b4f-4fce-8947-e8441b4a90af",)
+
+    def describe(self, file: Resource) -> Optional[FileDescription]: ...
+
+    def format(self, file: Resource) -> tuple: ...
+
+    def get_files(
+        self,
+        group: Union[List[str], str],
+        uf: Optional[Union[List[str], str]] = None,
+        year: Optional[Union[list, str, int]] = None,
+        month: Optional[Union[list, str, int]] = None,
+    ) -> List[Resource]: ...
+
+
+class SINAN(Dataset):
+    name = "SINAN"
+    ids = (
+        "4d5e5d44-58a8-4d67-b8aa-4ef1e4b00a1c",
+        "5699abe0-0510-4da8-b47d-209b3bb32b34",
+        "4557ba96-7d52-4a56-bd6f-f99a5af09f77",
+        "740ce8f4-7a5d-4351-aad4-7623f2490ada",
+    )
+
+    def describe(self, file: Resource) -> Optional[FileDescription]: ...
+
+    def format(self, file: Resource) -> tuple: ...
+
+    def get_files(
+        self,
+        dis_code: Optional[Union[str, list]] = None,
+        year: Optional[Union[str, int, list]] = None,
+    ) -> List[Resource]: ...