Skip to content

🧹 refactor: reduce parameters in _enrich method by grouping file attributes #33

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
devdudumuniz wants to merge 1 commit into
base: main
Choose a base branch
from
+75 −48
Draft

🧹 refactor: reduce parameters in _enrich method by grouping file attributes #33

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 50 additions & 29 deletions pysus/management/normalize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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(
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
44 changes: 25 additions & 19 deletions pysus/tests/management/test_normalize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
BucketNormalizer,
CatalogPathFix,
CatalogRowDelete,
FileAttributes,
ObjectRename,
formatter_for,
)
Expand Down Expand Up @@ -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:
Expand Down