diff --git a/pysus/management/normalize.py b/pysus/management/normalize.py index eec9a1f6..2634b56b 100644 --- a/pysus/management/normalize.py +++ b/pysus/management/normalize.py @@ -58,6 +58,14 @@ class CatalogPathFix: new_path: str +@dataclass +class FileAttributes: + group: str | None + year: int | None + month: int | None + state: str | None + + @dataclass class CatalogRowDelete: catalog: str @@ -448,23 +456,20 @@ def _enrich( origin: str, dataset: str, name: str, - group: str | None, - year: int | None, - month: int | None, - state: str | None, - ) -> dict: + attrs: FileAttributes, + ) -> FileAttributes: """Fill attribute gaps using the dataset formatter. Catalog values win; formatter output fills missing values and replaces legacy directory names (e.g. group ``"Dados"``) with the parsed group code. """ - enriched = { - "group": group, - "year": year, - "month": month, - "state": state, - } + enriched = FileAttributes( + group=attrs.group, + year=attrs.year, + month=attrs.month, + state=attrs.state, + ) formatter = formatter_for(origin, dataset) if formatter is None: return enriched @@ -477,16 +482,16 @@ def _enrich( if parsed_group and isinstance(parsed_group, dict): parsed_group = parsed_group.get("name") - if parsed_group and str(parsed_group) != enriched["group"]: + if parsed_group and str(parsed_group) != enriched.group: # formatters are curated; catalog groups may be legacy # (e.g. directory names) or NULL - enriched["group"] = str(parsed_group) - if enriched["year"] is None and parsed.get("year"): - enriched["year"] = int(parsed["year"]) - if enriched["month"] is None and parsed.get("month"): - enriched["month"] = int(parsed["month"]) - if enriched["state"] is None and parsed.get("state"): - enriched["state"] = str(parsed["state"]) + enriched.group = str(parsed_group) + if enriched.year is None and parsed.get("year"): + enriched.year = int(parsed["year"]) + if enriched.month is None and parsed.get("month"): + enriched.month = int(parsed["month"]) + if enriched.state is None and parsed.get("state"): + enriched.state = str(parsed["state"]) return enriched def survey_relayout( @@ -522,16 +527,24 @@ def survey_relayout( Path(origin_path).name if origin_path else Path(path).name ) enriched = self._enrich( - origin, dataset, source_name, group, year, month, state + origin, + dataset, + source_name, + FileAttributes( + group=group, + year=year, + month=month, + state=state, + ), ) new_key = compose_s3_key( origin=origin, dataset=dataset, name=source_name, - group=enriched["group"], - year=enriched["year"], - month=enriched["month"], - state=enriched["state"], + group=enriched.group, + year=enriched.year, + month=enriched.month, + state=enriched.state, ) if new_key == path: continue @@ -605,16 +618,24 @@ def relocate_uncataloged( continue name = Path(key).name enriched = self._enrich( - origin, dataset, name, None, None, None, None + origin, + dataset, + name, + FileAttributes( + group=None, + year=None, + month=None, + state=None, + ), ) new_key = compose_s3_key( origin=origin, dataset=dataset, name=name, - group=enriched["group"], - year=enriched["year"], - month=enriched["month"], - state=enriched["state"], + group=enriched.group, + year=enriched.year, + month=enriched.month, + state=enriched.state, ) if new_key == key: plan.raw_objects.append(key) diff --git a/pysus/tests/management/test_normalize.py b/pysus/tests/management/test_normalize.py index 35bcc21b..fbd0ce8a 100644 --- a/pysus/tests/management/test_normalize.py +++ b/pysus/tests/management/test_normalize.py @@ -8,6 +8,7 @@ BucketNormalizer, CatalogPathFix, CatalogRowDelete, + FileAttributes, ObjectRename, formatter_for, ) @@ -77,38 +78,43 @@ def test_short_key(self): class TestEnrich: def test_catalog_values_win(self, normalizer): enriched = normalizer._enrich( - "ftp", "SINAN", "DENGBR25.dbc", "DENG", 2025, None, None + "ftp", + "SINAN", + "DENGBR25.dbc", + FileAttributes("DENG", 2025, None, None), ) - assert enriched == { - "group": "DENG", - "year": 2025, - "month": None, - "state": None, - } + assert enriched.group == "DENG" + assert enriched.year == 2025 + assert enriched.month is None + assert enriched.state is None def test_formatter_fills_gaps(self, normalizer): enriched = normalizer._enrich( - "ftp", "SINAN", "DENGBR25.dbc", None, None, None, None + "ftp", + "SINAN", + "DENGBR25.dbc", + FileAttributes(None, None, None, None), ) - assert enriched["group"] == "DENG" - assert enriched["year"] == 2025 + assert enriched.group == "DENG" + assert enriched.year == 2025 def test_legacy_group_replaced_by_formatter(self, normalizer): enriched = normalizer._enrich( - "ftp", "CIHA", "CIHAMA2209.parquet", "Dados", 2022, 9, "MA" + "ftp", + "CIHA", + "CIHAMA2209.parquet", + FileAttributes("Dados", 2022, 9, "MA"), ) - assert enriched["group"] == "CIHA" + assert enriched.group == "CIHA" def test_unknown_formatter_keeps_values(self, normalizer): enriched = normalizer._enrich( - "ftp", "NOPE", "X.dbc", "G", 2020, 1, "AC" + "ftp", "NOPE", "X.dbc", FileAttributes("G", 2020, 1, "AC") ) - assert enriched == { - "group": "G", - "year": 2020, - "month": 1, - "state": "AC", - } + assert enriched.group == "G" + assert enriched.year == 2020 + assert enriched.month == 1 + assert enriched.state == "AC" class TestSurveyRelayout: