refactor(ducklake): implement abstract classes on the ducklake api classes

Author: luabida

pysus/api/dadosgov/databases.py

@@ -19,15 +19,18 @@ class CNES(Dataset):
         "9455b341-b06e-408e-8e10-54b32b3d74ec",
     )
 
-    def describe(self, file: Resource): ...
+    def describe(self, file: Resource):
+        ...
 
-    def format(self, file: Resource) -> tuple: ...
+    def format(self, file: Resource) -> tuple:
+        ...
 
     def get_files(
         self,
         year: Optional[Union[list, str, int]] = None,
         month: Optional[Union[list, str, int]] = None,
-    ) -> List[Resource]: ...
+    ) -> List[Resource]:
+        ...
 
 
 class PNI(Dataset):
@@ -42,32 +45,38 @@ class PNI(Dataset):
         "9a25b796-80e3-444a-a4e7-405f5596d8ab",
     )
 
-    def describe(self, file: Resource): ...
+    def describe(self, file: Resource):
+        ...
 
-    def format(self, file: Resource) -> tuple: ...
+    def format(self, file: Resource) -> tuple:
+        ...
 
     def get_files(
         self,
         year: Optional[Union[list, str, int]] = None,
         month: Optional[Union[list, str, int]] = None,
-    ) -> List[Resource]: ...
+    ) -> List[Resource]:
+        ...
 
 
 class SIA(Dataset):
     name = "SIA"
     ids = ("9a335cb7-2b4f-4fce-8947-e8441b4a90af",)
 
-    def describe(self, file: Resource): ...
+    def describe(self, file: Resource):
+        ...
 
-    def format(self, file: Resource) -> tuple: ...
+    def format(self, file: Resource) -> tuple:
+        ...
 
     def get_files(
         self,
         group: Union[List[str], str],
         uf: Optional[Union[List[str], str]] = None,
         year: Optional[Union[list, str, int]] = None,
         month: Optional[Union[list, str, int]] = None,
-    ) -> List[Resource]: ...
+    ) -> List[Resource]:
+        ...
 
 
 class SINAN(Dataset):
@@ -79,12 +88,15 @@ class SINAN(Dataset):
         "740ce8f4-7a5d-4351-aad4-7623f2490ada",
     )
 
-    def describe(self, file: Resource): ...
+    def describe(self, file: Resource):
+        ...
 
-    def format(self, file: Resource) -> tuple: ...
+    def format(self, file: Resource) -> tuple:
+        ...
 
     def get_files(
         self,
         dis_code: Optional[Union[str, list]] = None,
         year: Optional[Union[str, int, list]] = None,
-    ) -> List[Resource]: ...
+    ) -> List[Resource]:
+        ...

pysus/api/dadosgov/models.py

@@ -1,15 +1,12 @@
-import zipfile
 import pathlib
-import httpx
-import anyio
+import zipfile
 from datetime import datetime as dt
 from typing import Annotated, Any, List, Optional
 
-from pydantic import BaseModel, BeforeValidator, Field, ConfigDict
-from pysus.api.models import (
-    BaseRemoteFile,
-    BaseRemoteDataset,
-)
+import anyio
+import httpx
+from pydantic import BaseModel, BeforeValidator, ConfigDict, Field
+from pysus.api.models import BaseRemoteDataset, BaseRemoteFile
 
 
 def to_datetime(value: Any) -> Optional[dt]:
@@ -112,7 +109,9 @@ def describe(self, resource: Resource):
             group=self.slug,
             year=0,
             size=resource.api_size,
-            last_update=resource.last_modified or self.file_updated or dt.now(),
+            last_update=resource.last_modified
+            or self.file_updated
+            or dt.now(),
             uf=None,
             month=None,
             disease=self.title,

pysus/api/ducklake/catalog.py

@@ -9,9 +9,9 @@
     Index,
     Integer,
     String,
+    Table,
 )
 from sqlalchemy.orm import declarative_base, relationship
-from sqlalchemy import Table
 
 Base = declarative_base()
 

pysus/api/ducklake/client.py

@@ -1,22 +1,19 @@
-from typing import Optional, Any, List
 from pathlib import Path
-import httpx
-import duckdb
+from typing import Any, List, Optional
 
 import anyio
 import boto3
+import duckdb
+import httpx
 from botocore.config import Config
+from pydantic import BaseModel, PrivateAttr, SecretStr
+from pysus import CACHEPATH
+from pysus.api.models import BaseLocalFile, BaseRemoteClient
 from sqlalchemy import create_engine
-from sqlalchemy.orm import sessionmaker, joinedload
-from pydantic import PrivateAttr, SecretStr, BaseModel
+from sqlalchemy.orm import joinedload, sessionmaker
 
-from pysus import CACHEPATH
-from pysus.api.models import (
-    BaseLocalFile,
-    BaseRemoteClient,
-)
-from .models import CatalogDataset, CatalogFile
 from .catalog import Dataset, DatasetGroup
+from .models import CatalogDataset, CatalogFile
 
 
 class DuckLakeCredentials(BaseModel):
@@ -63,7 +60,8 @@ def _fetch():
                     .options(
                         joinedload(Dataset.dataset_metadata),
                         joinedload(Dataset.groups).joinedload(
-                            DatasetGroup.files),
+                            DatasetGroup.files
+                        ),
                     )
                     .all()
                 )
@@ -90,7 +88,9 @@ async def login(
         self._con = await anyio.to_thread.run_sync(self._create_connection)
 
         if self._is_authenticated:
-            self._s3_client = await anyio.to_thread.run_sync(self._get_s3_client)
+            self._s3_client = await anyio.to_thread.run_sync(
+                self._get_s3_client
+            )
 
     async def connect(self, force: bool = False):
         if self._con and not force:
@@ -191,23 +191,28 @@ def _create_connection(self):
         }
 
         if self._is_authenticated:
-            s3_cfg["s3_access_key_id"] = self.credentials.access_key.get_secret_value()
-            s3_cfg["s3_secret_access_key"] = (
-                self.credentials.secret_key.get_secret_value()
-            )
+            s3_cfg[
+                "s3_access_key_id"
+            ] = self.credentials.access_key.get_secret_value()
+            s3_cfg[
+                "s3_secret_access_key"
+            ] = self.credentials.secret_key.get_secret_value()
 
         for key, value in s3_cfg.items():
             con.execute(f"SET {key}='{value}';")
 
         mode = "" if self._is_authenticated else "(READ_ONLY)"
-        con.execute(f"ATTACH 'ducklake:{self._catalog_local}' AS pysus {mode};")
+        con.execute(
+            f"ATTACH 'ducklake:{self._catalog_local}' AS pysus {mode};"
+        )
         con.execute("USE pysus;")
         return con
 
     async def _upload_catalog(self):
         if not self._is_authenticated:
             raise PermissionError(
-                "Admin credentials required to upload catalog.")
+                "Admin credentials required to upload catalog."
+            )
 
         def _upload():
             self._s3_client.upload_file(

pysus/api/ducklake/models.py

@@ -1,18 +1,18 @@
 import hashlib
-from typing import List, Optional, Union
 from datetime import datetime
 from pathlib import Path
+from typing import List, Optional, Union
 
 import anyio
 from pydantic import Field
-
 from pysus.api.models import (
     BaseRemoteClient,
-    BaseRemoteGroup,
-    BaseRemoteFile,
     BaseRemoteDataset,
+    BaseRemoteFile,
+    BaseRemoteGroup,
 )
-from .catalog import File, Dataset, DatasetGroup
+
+from .catalog import Dataset, DatasetGroup, File
 
 
 class CatalogFile(BaseRemoteFile):
@@ -86,7 +86,9 @@ def long_name(self) -> str:
     @property
     def description(self) -> str:
         return (
-            self.record.group_metadata.description if self.record.group_metadata else ""
+            self.record.group_metadata.description
+            if self.record.group_metadata
+            else ""
         )
 
     async def files(self, **kwargs) -> List[CatalogFile]:
@@ -118,7 +120,9 @@ def description(self) -> str:
         )
 
     async def groups(self) -> List[CatalogGroup]:
-        return [CatalogGroup(record=g, dataset=self) for g in self.record.groups]
+        return [
+            CatalogGroup(record=g, dataset=self) for g in self.record.groups
+        ]
 
     async def files(self, **kwargs) -> List[CatalogFile]:
         return [CatalogFile(record=f, parent=self) for f in self.record.files]

pysus/api/extensions.py

@@ -8,25 +8,21 @@
 from pathlib import Path
 from typing import AsyncGenerator, Dict, List, Optional, Type, Union
 
-from pydantic import Field
 import anyio
 import chardet
 import fastparquet
 import magic
 import pandas as pd
 import pyarrow.parquet as pq
-
+from pydantic import Field
 from pysus import CACHEPATH
-from pysus.api.models import (
-    BaseCompressedFile,
-    BaseLocalFile,
-    BaseTabularFile,
-)
+from pysus.api.models import BaseCompressedFile, BaseLocalFile, BaseTabularFile
+
 from .types import FileType
 
 try:
-    from pyreaddbc import read_dbc, dbc2dbf
     from dbfread import DBF as DBFReader
+    from pyreaddbc import dbc2dbf, read_dbc
 
     FTP_IMPORT = True
 except ImportError:

pysus/api/ftp/models.py

@@ -9,17 +9,13 @@
 from ftplib import FTP
 from typing import Any, Dict, List, Optional, Tuple, TypedDict, Union
 
-from typing_extensions import Self
 from aioftp import Client
 from loguru import logger
-from tqdm import tqdm
-
 from pysus import CACHEPATH
-from pysus.api.models import (
-    BaseRemoteFile,
-    BaseRemoteDataset,
-)
+from pysus.api.models import BaseRemoteDataset, BaseRemoteFile
 from pysus.utils import to_list
+from tqdm import tqdm
+from typing_extensions import Self
 
 from .client import FTPSingleton
 
@@ -274,13 +270,17 @@ def load_directory_content(path: str) -> FileContent:
         def line_parser(line: str):
             if "<DIR>" in line:
                 date, time, _, name = line.strip().split(maxsplit=3)
-                modify = datetime.strptime(f"{date} {time}", "%m-%d-%y %I:%M%p")
+                modify = datetime.strptime(
+                    f"{date} {time}", "%m-%d-%y %I:%M%p"
+                )
                 info = {"size": 0, "type": "dir", "modify": modify}
                 xpath = f"{path}/{name}"
                 content[name] = Directory(xpath)
             else:
                 date, time, size, name = line.strip().split(maxsplit=3)
-                modify = datetime.strptime(f"{date} {time}", "%m-%d-%y %I:%M%p")
+                modify = datetime.strptime(
+                    f"{date} {time}", "%m-%d-%y %I:%M%p"
+                )
                 info: FileInfo = {
                     "size": size,
                     "type": "file",
@@ -353,7 +353,8 @@ def content(self) -> List[Union[Directory, FTPFile]]:
         """
         if not self.__content__:
             logger.info(
-                "content is not loaded, use `load()` to load default paths")
+                "content is not loaded, use `load()` to load default paths"
+            )
             return []
         return sorted(list(self.__content__.values()), key=str)
 
@@ -418,7 +419,9 @@ def get_files(self, *args, **kwargs) -> list[FTPFile]:
         """
         ...
 
-    def download(self, files: List[FTPFile], local_dir: str = CACHEPATH) -> List[str]:
+    def download(
+        self, files: List[FTPFile], local_dir: str = CACHEPATH
+    ) -> List[str]:
         """
         Downloads a list of Files.
         """
@@ -433,7 +436,9 @@ def download(self, files: List[FTPFile], local_dir: str = CACHEPATH) -> List[str
             return dfiles[0]
         return dfiles
 
-    async def async_download(self, files: List[FTPFile], local_dir: str = CACHEPATH):
+    async def async_download(
+        self, files: List[FTPFile], local_dir: str = CACHEPATH
+    ):
         """
         Asynchronously downloads a list of files
         """