Skip to content

Commit e0e4b16

Browse files
luabidaluabida
committed
fix: include group on sinan listing
1 parent c0d4296 commit e0e4b16

7 files changed

Lines changed: 210 additions & 242 deletions

File tree

‎README.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ The easiest way to get data as a pandas DataFrame:
3737
from pysus import sinan, sinasc, sim, sih, sia, pni, ibge, cnes, ciha
3838

3939
# Download SINAN Dengue data for 2024
40-
df = sinan(disease="deng", year=2024)
40+
df = sinan(disease="deng", year=2000)
4141

4242
# Multiple years
4343
df = sinan(disease="deng", year=[2023, 2024])
@@ -81,7 +81,7 @@ async def main():
8181
df = pysus.read_parquet(paths, mode="union").df()
8282
```
8383

84-
### Using the TUI
84+
### Using the TUI (unstable/under testing)
8585

8686
Launch the interactive text-based interface:
8787

‎pysus/api/_impl/databases.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -267,6 +267,7 @@ def _query():
267267
"name": r.path.split("/")[-1],
268268
"path": r.path,
269269
"dataset": r.dataset.name if r.dataset else None,
270+
"group": r.group.name if r.group else None,
270271
"year": r.year,
271272
"month": r.month,
272273
"state": r.state,

‎pysus/api/ducklake/client.py‎

Lines changed: 46 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,42 @@
1414
from sqlalchemy.pool import StaticPool
1515

1616
from .catalog import CatalogDataset, CatalogFile, DatasetGroup
17-
from .models import Dataset, File
17+
from .models import DuckDataset, File
18+
19+
20+
class CatalogDatasetAdapter:
21+
def __init__(self, catalog_dataset: CatalogDataset, ducklake):
22+
self.name = catalog_dataset.name
23+
self.long_name = catalog_dataset.long_name or ""
24+
self.description = catalog_dataset.description or ""
25+
self.group_definitions: dict[str, str] = {}
26+
self.ducklake = ducklake
27+
self.client = ducklake
28+
29+
@property
30+
def content(self):
31+
return self.ducklake.query(dataset=self.name.upper())
32+
33+
34+
class DatasetGroupAdapter:
35+
def __init__(self, dataset_group: DatasetGroup, dataset):
36+
self.name = dataset_group.name
37+
self.long_name = dataset_group.long_name or ""
38+
self.description = dataset_group.description or ""
39+
self.dataset = dataset
40+
41+
def __str__(self):
42+
return self.name
43+
44+
@property
45+
async def files(self):
46+
return []
47+
48+
async def _fetch_files(self):
49+
return []
50+
51+
async def search(self, **kwargs):
52+
return []
1853

1954

2055
class DuckLakeCredentials(BaseModel):
@@ -66,7 +101,7 @@ def _catalog_url(self) -> str:
66101
def _is_authenticated(self) -> bool:
67102
return self.credentials is not None
68103

69-
async def datasets(self, **kwargs) -> list[Dataset]:
104+
async def datasets(self, **kwargs) -> list[DuckDataset]:
70105
if not self._Session:
71106
await self.connect()
72107

@@ -86,7 +121,7 @@ def _fetch():
86121
return results
87122

88123
records = await to_thread.run_sync(_fetch)
89-
return [Dataset(record=rec, client=self) for rec in records]
124+
return [DuckDataset(record=rec, client=self) for rec in records]
90125

91126
async def login(
92127
self,
@@ -300,6 +335,13 @@ def _query():
300335

301336
records = await to_thread.run_sync(_query)
302337
return [
303-
File(path=r.path, record=r, dataset=r.dataset, group=r.group)
338+
File(
339+
path=r.path,
340+
record=r,
341+
dataset=CatalogDatasetAdapter(r.dataset, self),
342+
group=(
343+
DatasetGroupAdapter(r.group, r.dataset) if r.group else None
344+
),
345+
)
304346
for r in records
305347
]

‎pysus/api/ducklake/models.py‎

Lines changed: 44 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
from collections.abc import Callable
33
from datetime import datetime
44
from pathlib import Path
5+
from typing import Any, Union
56

67
import anyio
78
from pydantic import Field
@@ -19,6 +20,8 @@
1920
class File(BaseRemoteFile):
2021
record: CatalogFile = Field(exclude=True)
2122
type: str = "remote"
23+
dataset: Any
24+
group: Any = None
2225

2326
@property
2427
def basename(self) -> str:
@@ -51,6 +54,7 @@ async def _download(
5154
) -> Path:
5255
if not output:
5356
output = CACHEPATH / self.name
57+
5458
return await self.client._download_file(
5559
self,
5660
output,
@@ -72,41 +76,7 @@ def _calculate():
7276
return actual_hash == self.sha256
7377

7478

75-
class Group(BaseRemoteGroup):
76-
record: DatasetGroup = Field(exclude=True)
77-
dataset: "Dataset" = Field(exclude=True)
78-
79-
@property
80-
def name(self) -> str:
81-
return self.record.name
82-
83-
@property
84-
def long_name(self) -> str:
85-
return (
86-
self.record.group_metadata.long_name
87-
if self.record.group_metadata
88-
else self.name
89-
)
90-
91-
@property
92-
def description(self) -> str:
93-
if self.record.group_metadata:
94-
return self.record.group_metadata.description
95-
return ""
96-
97-
async def _fetch_files(self) -> list[BaseRemoteFile]:
98-
return [
99-
File(
100-
path=f.path,
101-
record=f,
102-
group=self,
103-
dataset=self.dataset,
104-
)
105-
for f in self.record.files
106-
]
107-
108-
109-
class Dataset(BaseRemoteDataset):
79+
class DuckDataset(BaseRemoteDataset):
11080
record: CatalogDataset = Field(exclude=True)
11181
client: BaseRemoteClient = Field(exclude=True)
11282

@@ -133,14 +103,12 @@ def description(self) -> str:
133103
else ""
134104
)
135105

136-
async def _fetch_content(
137-
self,
138-
) -> list[Group | File]:
139-
items: list[Group | File] = []
106+
async def _fetch_content(self) -> list[Union["DuckGroup", File]]:
107+
items: list[Union["DuckGroup", File]] = []
140108

141109
if self.record.groups:
142110
items.extend(
143-
[Group(record=g, dataset=self) for g in self.record.groups],
111+
[DuckGroup(record=g, dataset=self) for g in self.record.groups]
144112
)
145113

146114
if self.record.files:
@@ -152,7 +120,42 @@ async def _fetch_content(
152120
dataset=self,
153121
)
154122
for f in self.record.files
155-
],
123+
]
156124
)
157125

158126
return items
127+
128+
129+
class DuckGroup(BaseRemoteGroup):
130+
record: DatasetGroup = Field(exclude=True)
131+
dataset: DuckDataset = Field(exclude=True)
132+
133+
@property
134+
def name(self) -> str:
135+
return self.record.name
136+
137+
@property
138+
def long_name(self) -> str:
139+
return (
140+
self.record.group_metadata.long_name
141+
if self.record.group_metadata
142+
else self.name
143+
)
144+
145+
@property
146+
def description(self) -> str:
147+
if self.record.group_metadata:
148+
return self.record.group_metadata.description
149+
return ""
150+
151+
async def _fetch_files(self) -> list[BaseRemoteFile]:
152+
files: list[BaseRemoteFile] = [
153+
File(
154+
path=f.path,
155+
record=f,
156+
group=self,
157+
dataset=self.dataset,
158+
)
159+
for f in self.record.files
160+
]
161+
return files

‎pysus/api/ftp/databases.py‎

Lines changed: 46 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -403,12 +403,55 @@ class SINAN(Dataset):
403403
]
404404

405405
group_definitions: dict[str, str] = {
406-
"DENG": "Dengue",
407-
"ZIKA": "Zika Vírus",
406+
"ACBI": "Acidente de trabalho com material biológico",
407+
"ACGR": "Acidente de trabalho",
408+
"ANIM": "Acidente por Animais Peçonhentos",
409+
"ANTR": "Atendimento Antirrabico",
410+
"BOTU": "Botulismo",
411+
"CANC": "Cancêr relacionado ao trabalho",
412+
"CHAG": "Doença de Chagas Aguda",
408413
"CHIK": "Febre de Chikungunya",
414+
"COLE": "Cólera",
415+
"COQU": "Coqueluche",
416+
"DENG": "Dengue",
417+
"DERM": "Dermatoses ocupacionais",
418+
"DIFT": "Difteria",
419+
"ESQU": "Esquistossomose",
420+
"EXAN": "Doença exantemáticas",
421+
"FMAC": "Febre Maculosa",
422+
"FTIF": "Febre Tifóide",
409423
"HANS": "Hanseníase",
424+
"HANT": "Hantavirose",
425+
"HEPA": "Hepatites Virais",
426+
"IEXO": "Intoxicação Exógena",
427+
"INFL": "Influenza Pandêmica",
428+
"LEIV": "Leishmaniose Visceral",
429+
"LEPT": "Leptospirose",
430+
"LERD": "LER/Dort",
431+
"LTAN": "Leishmaniose Tegumentar Americana",
432+
"MALA": "Malária",
433+
"MENI": "Meningite",
434+
"MENT": "Transtornos mentais relacionados ao trabalho",
435+
"NTRA": "Notificação de Tracoma",
436+
"PAIR": "Perda auditiva por ruído relacionado ao trabalho",
437+
"PEST": "Peste",
438+
"PFAN": "Paralisia Flácida Aguda",
439+
"PNEU": "Pneumoconioses realacionadas ao trabalho",
440+
"RAIV": "Raiva",
441+
"SDTA": "Surto Doenças Transmitidas por Alimentos",
442+
"SIFA": "Sífilis Adquirida",
443+
"SIFC": "Sífilis Congênita",
444+
"SIFG": "Sífilis em Gestante",
445+
"SRC": "Síndrome da Rubéola Congênia",
446+
"TETA": "Tétano Acidental",
447+
"TETN": "Tétano Neonatal",
448+
"TOXC": "Toxoplasmose Congênita",
449+
"TOXG": "Toxoplasmose Gestacional",
450+
"TRAC": "Inquérito de Tracoma",
410451
"TUBE": "Tuberculose",
411-
"ANIM": "Acidente por Animais Peçonhentos",
452+
"VARC": "Varicela",
453+
"VIOL": "Violência doméstica, sexual e/ou outras violências",
454+
"ZIKA": "Zika Vírus",
412455
}
413456

414457
@property

‎pysus/api/models.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -327,6 +327,7 @@ async def search(self, **kwargs) -> list[BaseRemoteFile]:
327327

328328
class BaseRemoteDataset(BaseRemoteObject, SearchableMixin, ABC):
329329
client: BaseRemoteClient = Field(exclude=True)
330+
group_definitions: dict[str, str] = {}
330331
_content: Sequence[BaseRemoteGroup | BaseRemoteFile] | None = PrivateAttr(
331332
default=None
332333
)

0 commit comments

Comments
 (0)