🧹 Refactor ensure_group to take DatasetGroup object

pyproject.toml

@@ -59,6 +59,7 @@ pytest-retry = "1.7.0"
 datamodel-code-generator = {extras = ["http"], version = ">=0.54,<0.65"}
 pytest-asyncio = "^1.3.0"
 pytest-cov = "^7.1.0"
+pylint = "^4.0.8"
 
 [tool.poetry.group.docs.dependencies]
 sphinx = "^5.1.1"

pysus/management/catalog.py

@@ -15,6 +15,7 @@
 from typing import TYPE_CHECKING, Any
 
 from pysus.api.errors import CatalogError
+from pysus.api.metadata.models import DatasetGroup
 
 if TYPE_CHECKING:  # pragma: no cover
     from pysus.api.ducklake.client import DuckLake
@@ -154,19 +155,21 @@ def ensure_group(
         self,
         cursor,
         dataset_id: int,
-        name: str | None,
-        long_name: str | None = None,
-        description: str | None = None,
+        group: DatasetGroup | str | None,
     ) -> int | None:
         """Return the group id for (dataset, name), creating it if needed."""
-        if not name:
+        if not group:
             return None
 
-        name = name.strip().upper()
-        # ``long_name`` is NOT NULL; fall back to the group code itself when
-        # the source carries no long name (e.g. SIM's legacy "D" group).
-        if not long_name:
+        if isinstance(group, DatasetGroup):
+            name = group.name.strip().upper()
+            long_name = group.long_name or name
+            description = group.description
+        else:
+            name = str(group).strip().upper()
             long_name = name
+            description = None
+
         cursor.execute(
             "SELECT id, long_name, description FROM pysus.dataset_groups "
             "WHERE dataset_id = ? AND name = ?",

pysus/management/sync.py

@@ -29,6 +29,7 @@
 
 import anyio
 import httpx
+from pysus.api.metadata.models import DatasetGroup
 from pysus import CACHEPATH
 from pysus.api.ducklake.functional import upload_s3
 from pysus.api.errors import AuthenticationError, ConnectionError
@@ -1191,14 +1192,10 @@ def _catalog_rows(
         )
 
         group = getattr(file, "group", None)
-        group_name = getattr(group, "name", None) if group is not None else None
-        group_name = str(group_name) if group_name else None
         group_id = writer.ensure_group(
             dataset_cursor,
             dataset_id,
-            group_name,
-            getattr(group, "long_name", None) if group else None,
-            getattr(group, "description", None) if group else None,
+            group,
         )
 
         writer.upsert_file(
@@ -1436,8 +1433,11 @@ async def _fix_misparsed_metadata(
                     group_id = self.writer.ensure_group(
                         cursor,
                         dataset_id,
-                        str(group) if group else None,
-                        group_long_name or (str(group) if group else None),
+                        DatasetGroup(
+                            name=str(group),
+                            long_name=group_long_name or str(group),
+                            description="",
+                        ) if group else None,
                     )
                     cursor.execute(
                         "UPDATE pysus.files SET path = ?, month = ?, "

pysus/tests/management/test_catalog.py

@@ -7,6 +7,7 @@
 import pyarrow as pa
 import pytest
 from pysus.management.catalog import CatalogWriter
+from pysus.api.metadata.models import DatasetGroup
 
 _SCHEMA = """
 CREATE SCHEMA pysus;
@@ -492,7 +493,7 @@ def test_returns_none_for_empty_name(self, full_catalog):
     def test_creates_new_group(self, full_catalog):
         writer, cursor, _ = full_catalog
         gid = writer.ensure_group(
-            cursor, dataset_id=1, name="deng", long_name="Dengue"
+            cursor, dataset_id=1, group=DatasetGroup(name="deng", long_name="Dengue", description="")
         )
         assert gid == 1
         cursor.execute(
@@ -505,10 +506,10 @@ def test_creates_new_group(self, full_catalog):
     def test_reuses_existing_group(self, full_catalog):
         writer, cursor, _ = full_catalog
         gid1 = writer.ensure_group(
-            cursor, dataset_id=1, name="deng", long_name="Dengue"
+            cursor, dataset_id=1, group=DatasetGroup(name="deng", long_name="Dengue", description="")
         )
         gid2 = writer.ensure_group(
-            cursor, dataset_id=1, name="deng", long_name="Dengue v2"
+            cursor, dataset_id=1, group=DatasetGroup(name="deng", long_name="Dengue v2", description="")
         )
         assert gid1 == gid2
         cursor.execute(
@@ -519,7 +520,7 @@ def test_reuses_existing_group(self, full_catalog):
 
     def test_group_name_uppercased(self, full_catalog):
         writer, cursor, _ = full_catalog
-        gid = writer.ensure_group(cursor, dataset_id=1, name="  sinh  ")
+        gid = writer.ensure_group(cursor, dataset_id=1, group="  sinh  ")
         cursor.execute(
             "SELECT name FROM pysus.dataset_groups WHERE id = ?", (gid,)
         )