tests: include more tests

Author: luabida

pysus/api/client.py

@@ -26,7 +26,7 @@
 from .ftp import FTPClient
 from .models import BaseLocalFile, BaseRemoteFile
 
-if TYPE_CHECKING:
+if TYPE_CHECKING:  # pragma: no cover
     from duckdb import DuckDBPyConnection
 
 
@@ -403,7 +403,7 @@ async def download_to_parquet(
 
         if hasattr(local_file, "to_parquet"):
             original_path = local_file.path
-            parquet_file = await local_file.to_parquet(callback=callback)  # type: ignore
+            parquet_file = await local_file.to_parquet(callback=callback)
             parquet_file.add_dv = add_dv
 
             await self._update_state(

pysus/api/dadosgov/client.py

@@ -13,7 +13,7 @@
 from pysus.api.models import BaseRemoteClient, BaseRemoteFile
 from pysus.api.types import DADOSGOV
 
-if TYPE_CHECKING:
+if TYPE_CHECKING:  # pragma: no cover
     from .models import Dataset
 
 

pysus/api/dadosgov/databases.py

@@ -272,7 +272,7 @@ def formatter(self, filename: str) -> dict[str, Any]:
                 }
 
             m = re.search(r"_(\w{3})-out_(\d{4})_\.csv$", name)
-            if m:
+            if m:  # pragma: no cover
                 return {
                     "state": None,
                     "year": _parse_year(m.group(2)),

pysus/api/ducklake/catalog/orm/dataset.py

@@ -74,12 +74,12 @@ class Dataset(Base):
     description = Column(String, nullable=True)
 
     groups = relationship(
-        "DatasetGroup",
+        "Group",
         back_populates="dataset",
         cascade="all, delete-orphan",
     )
     files = relationship(
-        "CatalogFile",
+        "File",
         back_populates="dataset",
         cascade="all, delete-orphan",
     )
@@ -128,9 +128,9 @@ class ColumnDefinition(Base):
     description = Column(String, nullable=True)
     nullable = Column(Boolean, nullable=False, default=True)
 
-    dataset = relationship("CatalogDataset", back_populates="columns")
+    dataset = relationship("Dataset", back_populates="columns")
     files = relationship(
-        "CatalogFile",
+        "File",
         secondary=file_columns,
         back_populates="columns",
     )
@@ -177,11 +177,11 @@ class Group(Base):
     description = Column(String, nullable=True)
 
     dataset = relationship(
-        "CatalogDataset",
+        "Dataset",
         back_populates="groups",
     )
     files = relationship(
-        "CatalogFile",
+        "File",
         back_populates="group",
         cascade="all, delete-orphan",
     )
@@ -281,11 +281,11 @@ class File(Base):
     )
 
     dataset: Mapped["Dataset"] = relationship(
-        "CatalogDataset",
+        "Dataset",
         back_populates="files",
     )
     group: Mapped[Optional["Group"]] = relationship(
-        "DatasetGroup",
+        "Group",
         back_populates="files",
     )
     columns: Mapped[list["ColumnDefinition"]] = relationship(

pysus/api/ducklake/catalog/orm/default.py

@@ -15,7 +15,8 @@ class Base(DeclarativeBase):
 
 
 class Dataset(Base):
-    """ORM model for the datasets table — central registry of available datasets.
+    """ORM model for the datasets table — central registry of available
+    datasets.
 
     Parameters
     ----------

pysus/api/ducklake/client.py

@@ -77,8 +77,8 @@ def __init__(self, engine=None, **data) -> None:
         self._engine = engine
         self._cache_dir: Path = Path(CACHEPATH) / "ducklake"
         self._cache_dir.mkdir(parents=True, exist_ok=True)
-        self._catalog_local: Path = self._cache_dir / "catalog.db"
-        self._catalog_remote: str = "public/catalog.db"
+        self._catalog_local: Path = self._cache_dir / "catalog.duckdb"
+        self._catalog_remote: str = "public/catalog.duckdb"
 
     @property
     def name(self) -> str:
@@ -197,7 +197,8 @@ def _setup_engine(self, local_path: Path | None = None):
         Parameters
         ----------
         local_path : Path, optional
-            Path to the catalog database file. Defaults to the discovery catalog.
+            Path to the catalog database file.
+            Defaults to the discovery catalog.
         """
         if local_path is None:
             local_path = self._catalog_local
@@ -210,7 +211,8 @@ def _setup_engine(self, local_path: Path | None = None):
             conn.exec_driver_sql("INSTALL ducklake; LOAD ducklake;")
 
             has_pysus = conn.exec_driver_sql(
-                "SELECT 1 FROM information_schema.schemata WHERE schema_name = 'pysus'"
+                "SELECT 1 FROM information_schema.schemata"
+                " WHERE schema_name = 'pysus'"
             ).fetchone()
 
             if has_pysus:
@@ -326,7 +328,9 @@ async def _download(
                 else:
                     raise e
 
-    async def _download_catalog(self, local_path: Path, remote_path: str) -> None:
+    async def _download_catalog(
+        self, local_path: Path, remote_path: str
+    ) -> None:
         """Download a catalog database from remote storage with retries.
 
         Parameters
@@ -380,7 +384,9 @@ def _get_s3_client(self):
             "s3",
             endpoint_url=f"https://{self.endpoint}",
             aws_access_key_id=self.credentials.access_key.get_secret_value(),
-            aws_secret_access_key=(self.credentials.secret_key.get_secret_value()),
+            aws_secret_access_key=(
+                self.credentials.secret_key.get_secret_value()
+            ),
             region_name=self.region,
             config=Config(signature_version="s3v4"),
         )
@@ -400,7 +406,10 @@ async def _upload_catalog(self) -> None:
             if not ds._catalog_local.exists():
                 continue
 
-            def _upload(local=str(ds._catalog_local), name=ds._catalog_name):
+            _local = str(ds._catalog_local)
+            _name = ds._catalog_name
+
+            def _upload(local=_local, name=_name):
                 self._s3_client.upload_file(
                     local,
                     self.bucket,

pysus/api/ducklake/models.py

@@ -13,16 +13,13 @@
 from anyio import to_thread
 from pydantic import Field, PrivateAttr
 from pysus import CACHEPATH
+from pysus.api.ducklake.catalog.orm.dataset import Dataset
+from pysus.api.ducklake.catalog.orm.dataset import File as CatalogFile
+from pysus.api.ducklake.catalog.orm.dataset import Group
 from pysus.api.models import BaseRemoteDataset, BaseRemoteFile, BaseRemoteGroup
 from sqlalchemy.orm import contains_eager, joinedload, sessionmaker
 
-from pysus.api.ducklake.catalog.orm.dataset import (
-    Dataset,
-    File as CatalogFile,
-    Group,
-)
-
-if TYPE_CHECKING:
+if TYPE_CHECKING:  # pragma: no cover
     from .client import DuckLake
 
 
@@ -41,23 +38,19 @@ class File(BaseRemoteFile):
         The parent group object, if any.
     """
 
-    def __init__(
-        self,
-        dataset: "DuckDataset",
-        record: CatalogFile,
-        group: Optional["DuckGroup"] = None,
-    ) -> None:
+    record: CatalogFile = Field(exclude=True)
+    group: Optional["DuckGroup"] = Field(default=None, exclude=True)
+
+    def __init__(self, **data: Any) -> None:
+        record = data.pop("record")
+        group = data.pop("group", None)
         super().__init__(
             path=Path(record.path),
             type=record.type or "remote",
-            dataset=dataset,
+            record=record,  # type: ignore[call-arg]
+            group=group,
+            **data,
         )
-        self.record: CatalogFile = record
-        self.group: Optional["DuckGroup"] = group
-
-    @property
-    def path(self) -> Path:
-        return Path(self.record.path)
 
     @property
     def basename(self) -> str:
@@ -191,7 +184,7 @@ def __init__(self, **data) -> None:
         super().__init__(**data)
         self._cache_dir: Path = Path(CACHEPATH) / "ducklake"
         self._cache_dir.mkdir(parents=True, exist_ok=True)
-        self._catalog_name: str = f"catalog_{self.record.name.lower()}.db"
+        self._catalog_name: str = f"catalog_{self.record.name.lower()}.duckdb"
         self._catalog_local: Path = self._cache_dir / self._catalog_name
 
     def __repr__(self) -> str:
@@ -306,7 +299,7 @@ def _upload():
             self.client._s3_client.upload_file(
                 str(self._catalog_local),
                 self.client.bucket,
-                f"catalog_{self.record.name.lower()}.db",
+                f"catalog_{self.record.name.lower()}.duckdb",
             )
 
         await to_thread.run_sync(_upload)

pysus/api/extensions.py

@@ -10,7 +10,6 @@
 from collections.abc import AsyncGenerator, Callable
 from datetime import datetime
 from pathlib import Path
-from typing import ClassVar
 
 import chardet
 import pandas as pd
@@ -19,6 +18,7 @@
 from anyio import to_thread
 from dbfread import DBF as DBFReader
 from pydantic import Field, PrivateAttr
+from pyreaddbc import dbc2dbf
 from pysus import CACHEPATH
 from pysus.api.metadata.models import Column
 from pysus.api.models import BaseCompressedFile, BaseLocalFile, BaseTabularFile
@@ -60,14 +60,6 @@ def _map_dtype(raw: str) -> str:
     return "VARCHAR"
 
 
-try:
-    from pyreaddbc import dbc2dbf
-
-    DBC_IMPORT = True
-except ImportError:
-    DBC_IMPORT = False
-
-
 class File(BaseLocalFile):
     """Represents a generic local file with no special handling."""
 
@@ -839,76 +831,6 @@ def _extract():
         return list(await asyncio.gather(*tasks))
 
 
-class DBCNotImported(BaseTabularFile):
-    """Placeholder for DBC files when optional dependency is not installed."""
-
-    path: Path = Field(default_factory=lambda: Path("..."))
-    type: str | FileType = Field(default="remote")
-    import_err: ClassVar[
-        str
-    ] = """
-        run "pip install pysus[dbc]" to handle DBC files.
-        Make sure you also have libffi installed on the system. It may not work
-        on Windows
-    """
-
-    @property
-    def name(self) -> str:
-        """Raise ImportError indicating the missing DBC dependency."""
-        raise ImportError(self.import_err)
-
-    @property
-    def extension(self) -> str:
-        """Return the .dbc extension."""
-        return ".dbc"
-
-    @property
-    def size(self) -> int:
-        """Raise ImportError indicating the missing DBC dependency."""
-        raise ImportError(self.import_err)
-
-    @property
-    def modify(self) -> datetime:
-        """Raise ImportError indicating the missing DBC dependency."""
-        raise ImportError(self.import_err)
-
-    @property
-    def columns(self) -> list["Column"]:
-        """Raise ImportError indicating the missing DBC dependency."""
-        raise ImportError(self.import_err)
-
-    @property
-    def rows(self) -> int:
-        """Raise ImportError indicating the missing DBC dependency."""
-        raise ImportError(self.import_err)
-
-    async def load(self) -> pd.DataFrame:
-        """Raise ImportError indicating the missing DBC dependency."""
-        raise ImportError(self.import_err)
-
-    def stream(
-        self,
-        chunk_size: int = 10000,
-    ) -> AsyncGenerator[pd.DataFrame, None]:
-        """Raise ImportError indicating the missing DBC dependency."""
-
-        async def _internal_gen():
-            """Yield nothing; always raises ImportError."""
-            raise ImportError(self.import_err)
-            yield pd.DataFrame()
-
-        return _internal_gen()
-
-    async def to_parquet(
-        self,
-        output_path: str | Path | None = None,
-        chunk_size: int = 10000,
-        callback: Callable[[int, int], None] | None = None,
-    ) -> Parquet:
-        """Raise ImportError indicating the missing DBC dependency."""
-        raise ImportError(self.import_err)
-
-
 class ExtensionFactory:
     """Factory that maps file extensions and MIME types to handler classes."""
 
@@ -930,7 +852,7 @@ class ExtensionFactory:
         ".csv": CSV,
         ".parquet": Parquet,
         ".dbf": DBF,
-        ".dbc": DBC if DBC_IMPORT else DBCNotImported,  # type: ignore
+        ".dbc": DBC,
         ".pdf": PDF,
         ".json": JSON,
     }

pysus/api/ftp/client.py

@@ -13,7 +13,7 @@
 from pysus.api.models import BaseRemoteClient, BaseRemoteFile
 from pysus.api.types import FTP as FTP_STR
 
-if TYPE_CHECKING:
+if TYPE_CHECKING:  # pragma: no cover
     from pysus.api.ftp.models import Dataset
     from pysus.api.types import State
 

pysus/api/metadata/models.py

@@ -71,7 +71,8 @@ class Column:
 
     @classmethod
     def from_schema(cls, name: str, dtype: ColumnType) -> "Column":
-        """Create a Column from a file schema, looking up description from columns.py."""
+        """Create a Column from a file schema, looking up description from
+        columns.py."""
         return cls(
             name=name,
             description=pick_description(lookup_column_meta(name)),