sera-2 1.18.1__py3-none-any.whl → 1.19.0__py3-none-any.whl

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/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,59 @@ 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("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:
-                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)
+                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:
+                session.execute(
+                    text(
+                        f"SELECT setval('{tbl.__tablename__}_id_seq', (SELECT MAX(id) FROM \"{tbl.__tablename__}\"));"
+                    )
                 )
-            )
         session.commit()
 
 

{sera_2-1.18.1.dist-info → sera_2-1.19.0.dist-info}/METADATA

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

{sera_2-1.18.1.dist-info → sera_2-1.19.0.dist-info}/RECORD

@@ -26,9 +26,9 @@ sera/make/make_python_api.py,sha256=iXGbKQ3IJvsY1ur_fhurr_THFNnH66E3Wl85o0emUbw,
 sera/make/make_python_model.py,sha256=Nc4vDGgM8icgWBqzNnMgEkLadf5EsZwbbHs3WLW9_co,62778
 sera/make/make_python_services.py,sha256=0ZpWLwQ7Nwfn8BXAikAB4JRpNknpSJyJgY5b1cjtxV4,2073
 sera/make/make_typescript_model.py,sha256=1ouYFCeqOlwEzsGBiXUn4VZtLJjJW7GSacdOSlQzhjI,67012
-sera/misc/__init__.py,sha256=Tali_UBtwemETM30a6sP6BbwBMHr3hklPCX0bgiAcbw,513
+sera/misc/__init__.py,sha256=rOmGMv7QNzpSKZSyxChbRmEnBr3O443UlLGS0FIs3AI,561
 sera/misc/_formatter.py,sha256=aCGYL08l8f3aLODHxSocxBBwkRYEo3K1QzCDEn3suj0,1685
-sera/misc/_utils.py,sha256=f5mOgDlGh-OVwd6DXou2gTo9eRvJGK_aUT7pM3qzr98,4882
+sera/misc/_utils.py,sha256=O0Qh_BZXX7kmzhFBTnnxJT4BUxQbwNovhepSwhh33Ow,6532
 sera/models/__init__.py,sha256=vJC5Kzo_N7wd16ocNPy1VvAZDGNiWeiAhWJ4ihATKvA,780
 sera/models/_class.py,sha256=1J4Bd_LanzhhDWwZFHWGtFYD7lupe_alaB3D02ebNDI,2862
 sera/models/_collection.py,sha256=ZnQEriKC4X88Zz48Kn1AVZKH-1_l8OgWa-zf2kcQOOE,1414
@@ -42,6 +42,6 @@ sera/models/_parse.py,sha256=ciTLzCkO0q6xA1R_rHbnYJYK3Duo2oh56WeuwxXwJaI,12392
 sera/models/_property.py,sha256=9yMDxrmbyuF6-29lQjiq163Xzwbk75TlmGBpu0NLpkI,7485
 sera/models/_schema.py,sha256=VxJEiqgVvbXgcSUK4UW6JnRcggk4nsooVSE6MyXmfNY,1636
 sera/typing.py,sha256=o_DKfSvs8JpNRQ0kdaTc3BbfdkvibY3uY4tJRt-n2fQ,1023
-sera_2-1.18.1.dist-info/METADATA,sha256=TATTG19o7HW6O681m1dUdFs92SJ6oqizZUT_vF52zx8,936
-sera_2-1.18.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-sera_2-1.18.1.dist-info/RECORD,,
+sera_2-1.19.0.dist-info/METADATA,sha256=22HLRJ7LXZYcBduEokfXaTnePGNCUIyARqVDZFP404k,936
+sera_2-1.19.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+sera_2-1.19.0.dist-info/RECORD,,