sera-2 1.18.1__tar.gz → 1.19.1__tar.gz

{sera_2-1.18.1 → sera_2-1.19.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: sera-2
-Version: 1.18.1
+Version: 1.19.1
 Summary: 
 Author: Binh Vu
 Author-email: bvu687@gmail.com

{sera_2-1.18.1 → sera_2-1.19.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sera-2"
-version = "1.18.1"
+version = "1.19.1"
 description = ""
 authors = ["Binh Vu <bvu687@gmail.com>"]
 readme = "README.md"

{sera_2-1.18.1 → sera_2-1.19.1}/sera/make/make_python_model.py

@@ -720,6 +720,12 @@ def make_python_data_model(
 
         program.import_("__future__.annotations", True)
         program.import_("msgspec", False)
+
+        ident_manager = ImportHelper(
+            program,
+            GLOBAL_IDENTS,
+        )
+
         if cls.db is not None:
             # if the class is stored in the database, we need to import the database module
             program.import_(
@@ -824,7 +830,9 @@ def make_python_data_model(
                 lambda ast: ast.func(
                     "as_composite",
                     vars=as_composite_args,
-                    return_type=expr.ExprIdent(f"Optional[{cls.name}]"),
+                    return_type=PredefinedFn.item_getter(
+                        ident_manager.use("Optional"), expr.ExprIdent(cls.name)
+                    ),
                     comment="Create an embedded instance from the embedded columns in the database table. If all properties of this embedded class are None (indicating that the parent field is None), then this function will return None.",
                 )(
                     lambda ast_l1: ast_l1.if_(

{sera_2-1.18.1 → sera_2-1.19.1}/sera/misc/__init__.py

@@ -1,5 +1,6 @@
 from sera.misc._formatter import File, Formatter
 from sera.misc._utils import (
+    LoadTableDataArgs,
     assert_isinstance,
     assert_not_null,
     filter_duplication,
@@ -24,4 +25,5 @@ __all__ = [
     "load_data",
     "identity",
     "get_classpath",
+    "LoadTableDataArgs",
 ]

{sera_2-1.18.1 → sera_2-1.19.1}/sera/misc/_utils.py

@@ -1,10 +1,12 @@
 from __future__ import annotations
 
 import re
+from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Callable, Iterable, Optional, Type, TypeVar
+from typing import Any, Callable, Iterable, Optional, Sequence, Type, TypedDict, TypeVar
 
 import serde.csv
+import serde.json
 from sqlalchemy import Engine, text
 from sqlalchemy.orm import Session
 from tqdm import tqdm
@@ -113,11 +115,19 @@ def filter_duplication(
     return new_lst
 
 
+class LoadTableDataArgs(TypedDict, total=False):
+    table: type
+    tables: Sequence[type]
+    file: Path
+    files: Sequence[Path]
+    file_deser: Callable[[Path], list[dict]]
+    record_deser: Callable[[dict], Any | list[Any]]
+
+
 def load_data(
     engine: Engine,
     create_db_and_tables: Callable[[], None],
-    table_files: list[tuple[type, Path]],
-    table_desers: dict[type, Callable[[dict], Any]],
+    table_data: Sequence[LoadTableDataArgs],
     verbose: bool = False,
 ):
     """
@@ -133,23 +143,71 @@ def load_data(
     with Session(engine) as session:
         create_db_and_tables()
 
-        for tbl, file in tqdm(table_files, disable=not verbose, desc="Loading data"):
-            if file.name.endswith(".csv"):
-                records = serde.csv.deser(file, deser_as_record=True)
+        for args in tqdm(table_data, disable=not verbose, desc="Loading data"):
+            if "table" in args:
+                tbls = [args["table"]]
+            elif "tables" in args:
+                tbls = args["tables"]
+            else:
+                raise ValueError("Either 'table' or 'tables' must be provided in args.")
+
+            if "file" in args:
+                assert isinstance(args["file"], Path), "File must be a Path object."
+                files = [args["file"]]
+            elif "files" in args:
+                assert all(
+                    isinstance(f, Path) for f in args["files"]
+                ), "Files must be Path objects."
+                files = args["files"]
             else:
-                raise ValueError(f"Unsupported file format: {file.name}")
-            deser = table_desers[tbl]
-            records = [deser(row) for row in records]
-            for r in tqdm(records, desc=f"load {tbl.__name__}", disable=not verbose):
-                session.merge(r)
+                raise ValueError("Either 'file' or 'files' must be provided in args.")
+
+            raw_records = []
+            if "file_deser" not in args:
+                for file in files:
+                    if file.name.endswith(".csv"):
+                        raw_records.extend(serde.csv.deser(file, deser_as_record=True))
+                    elif file.name.endswith(".json"):
+                        raw_records.extend(serde.json.deser(file))
+                    else:
+                        raise ValueError(f"Unsupported file format: {file.name}")
+            else:
+                for file in files:
+                    raw_records.extend(args["file_deser"](file))
+
+            deser = args["record_deser"]
+            records = [deser(row) for row in raw_records]
+            for r in tqdm(
+                records,
+                desc=f"load {', '.join(tbl.__name__ for tbl in tbls)}",
+                disable=not verbose,
+            ):
+                if isinstance(r, Sequence):
+                    for x in r:
+                        session.merge(x)
+                else:
+                    session.merge(r)
             session.flush()
 
             # Reset the sequence for each table
-            session.execute(
-                text(
-                    f"SELECT setval('{tbl.__tablename__}_id_seq', (SELECT MAX(id) FROM \"{tbl.__tablename__}\"));"
+            for tbl in tbls:
+                # Check if the table has an auto-incrementing primary key
+                if not hasattr(tbl, "__table__") or not tbl.__table__.primary_key:
+                    continue
+
+                pk_columns = tbl.__table__.primary_key.columns
+                has_foreign_key = any(len(col.foreign_keys) > 0 for col in pk_columns)
+                has_auto_increment = any(
+                    col.autoincrement and col.type.python_type in (int,)
+                    for col in pk_columns
+                )
+                if has_foreign_key or not has_auto_increment:
+                    continue
+                session.execute(
+                    text(
+                        f"SELECT setval('{tbl.__tablename__}_id_seq', (SELECT MAX(id) FROM \"{tbl.__tablename__}\"));"
+                    )
                 )
-            )
         session.commit()
 
 

{sera_2-1.18.1 → sera_2-1.19.1}/sera/typing.py

@@ -38,4 +38,5 @@ GLOBAL_IDENTS = {
     "ASGIConnection": "litestar.connection.ASGIConnection",
     "UNSET": "sera.typing.UNSET",
     "ForeignKey": "sqlalchemy.ForeignKey",
+    "Optional": "typing.Optional",
 }