sera-2 1.2.0__py3-none-any.whl → 1.4.2__py3-none-any.whl

sera/make/make_python_model.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
-from operator import is_
-from typing import Sequence
+from typing import Callable, Sequence
 
 from codegen.models import AST, DeferredVar, PredefinedFn, Program, expr, stmt
 from sera.misc import (
@@ -16,12 +15,22 @@ from sera.models import (
     DataProperty,
     ObjectProperty,
     Package,
+    PyTypeWithDep,
     Schema,
 )
+from sera.typing import ObjectPath
 
 
-def make_python_data_model(schema: Schema, target_pkg: Package):
-    """Generate public classes for the API from the schema."""
+def make_python_data_model(
+    schema: Schema, target_pkg: Package, reference_classes: dict[str, ObjectPath]
+):
+    """Generate public classes for the API from the schema.
+
+    Args:
+        schema: The schema to generate the classes from.
+        target_pkg: The package to write the classes to.
+        reference_classes: A dictionary of class names to their references (e.g., the ones that are defined outside and used as referenced such as Tenant).
+    """
     app = target_pkg.app
 
     def from_db_type_conversion(
@@ -30,23 +39,29 @@ def make_python_data_model(schema: Schema, target_pkg: Package):
         value = PredefinedFn.attr_getter(record, expr.ExprIdent(prop.name))
         if isinstance(prop, ObjectProperty) and prop.target.db is not None:
             if prop.cardinality.is_star_to_many():
-                return PredefinedFn.map_list(
+                value = PredefinedFn.map_list(
                     value,
-                    lambda item: expr.ExprFuncCall(
-                        PredefinedFn.attr_getter(
-                            expr.ExprIdent(prop.target.name), expr.ExprIdent("from_db")
-                        ),
-                        [PredefinedFn.attr_getter(item, expr.ExprIdent(prop.name))],
+                    lambda item: PredefinedFn.attr_getter(
+                        item, expr.ExprIdent(prop.name + "_id")
                     ),
                 )
             else:
-                return expr.ExprFuncCall(
-                    PredefinedFn.attr_getter(
-                        expr.ExprIdent(prop.target.name), expr.ExprIdent("from_db")
-                    ),
-                    [value],
+                value = PredefinedFn.attr_getter(
+                    record, expr.ExprIdent(prop.name + "_id")
                 )
 
+            target_idprop = assert_not_null(prop.target.get_id_property())
+            conversion_fn = get_data_conversion(
+                target_idprop.datatype.get_python_type().type,
+                target_idprop.get_data_model_datatype().get_python_type().type,
+            )
+            value = conversion_fn(value)
+        elif isinstance(prop, DataProperty) and prop.is_diff_data_model_datatype():
+            value = get_data_conversion(
+                prop.datatype.get_python_type().type,
+                prop.get_data_model_datatype().get_python_type().type,
+            )(value)
+
         return value
 
     def to_db_type_conversion(
@@ -56,46 +71,75 @@ def make_python_data_model(schema: Schema, target_pkg: Package):
         prop: DataProperty | ObjectProperty,
     ):
         value = PredefinedFn.attr_getter(slf, expr.ExprIdent(prop.name))
-        if (
-            isinstance(prop, ObjectProperty)
-            and prop.target.db is not None
-            and prop.cardinality == Cardinality.MANY_TO_MANY
-        ):
-            # we have to use the associated object
-            AssociationTable = f"{cls.name}{prop.target.name}"
-            program.import_(
-                app.models.db.path
-                + f".{to_snake_case(AssociationTable)}.{AssociationTable}",
-                True,
-            )
-            return PredefinedFn.map_list(
-                value,
-                lambda item: expr.ExprFuncCall(
-                    expr.ExprIdent(AssociationTable),
-                    [PredefinedFn.keyword_assignment(f"{prop.name}_id", item)],
-                ),
-            )
+        if isinstance(prop, ObjectProperty):
+            if (
+                prop.target.db is not None
+                and prop.cardinality == Cardinality.MANY_TO_MANY
+            ):
+                # we have to use the associated object
+                # if this isn't a many-to-many relationship, we only keep the id, so no need to convert to the type.
+                AssociationTable = f"{cls.name}{prop.target.name}"
+                program.import_(
+                    app.models.db.path
+                    + f".{to_snake_case(AssociationTable)}.{AssociationTable}",
+                    True,
+                )
+
+                target_idprop = assert_not_null(prop.target.get_id_property())
+                conversion_fn = get_data_conversion(
+                    target_idprop.get_data_model_datatype().get_python_type().type,
+                    target_idprop.datatype.get_python_type().type,
+                )
+
+                return PredefinedFn.map_list(
+                    value,
+                    lambda item: expr.ExprFuncCall(
+                        expr.ExprIdent(AssociationTable),
+                        [
+                            PredefinedFn.keyword_assignment(
+                                f"{prop.name}_id", conversion_fn(item)
+                            )
+                        ],
+                    ),
+                )
+            elif prop.target.db is None:
+                # if the target class is not in the database, we need to convert the value to the python type used in db.
+                # if the cardinality is many-to-many, we need to convert each item in the list.
+                if prop.cardinality.is_star_to_many():
+                    value = PredefinedFn.map_list(
+                        value, lambda item: expr.ExprMethodCall(item, "to_db", [])
+                    )
+                else:
+                    value = expr.ExprMethodCall(value, "to_db", [])
+        elif isinstance(prop, DataProperty) and prop.is_diff_data_model_datatype():
+            # convert the value to the python type used in db.
+            value = get_data_conversion(
+                prop.get_data_model_datatype().get_python_type().type,
+                prop.datatype.get_python_type().type,
+            )(value)
         return value
 
-    def make_create(program: Program, cls: Class):
+    def make_upsert(program: Program, cls: Class):
         program.import_("__future__.annotations", True)
         program.import_("msgspec", False)
-        program.import_(
-            app.models.db.path + f".{cls.get_pymodule_name()}.{cls.name}",
-            True,
-            alias=f"{cls.name}DB",
-        )
+        if cls.db is not None:
+            # if the class is stored in the database, we need to import the database module
+            program.import_(
+                app.models.db.path + f".{cls.get_pymodule_name()}.{cls.name}",
+                True,
+                alias=f"{cls.name}DB",
+            )
         cls_ast = program.root.class_(
             "Upsert" + cls.name, [expr.ExprIdent("msgspec.Struct")]
         )
         for prop in cls.properties.values():
             # this is a create object, so users can create private field
             # hence, we do not check for prop.is_private
-            # if prop.is_private:
+            # if prop.data.is_private:
             #     continue
 
             if isinstance(prop, DataProperty):
-                pytype = prop.datatype.get_python_type()
+                pytype = prop.get_data_model_datatype().get_python_type()
                 if pytype.dep is not None:
                     program.import_(pytype.dep, True)
                 cls_ast(stmt.DefClassVarStatement(prop.name, pytype.type))
@@ -104,19 +148,21 @@ def make_python_data_model(schema: Schema, target_pkg: Package):
                     # if the target class is in the database, we expect the user to pass the foreign key for it.
                     pytype = (
                         assert_not_null(prop.target.get_id_property())
-                        .datatype.get_python_type()
-                        .type
+                        .get_data_model_datatype()
+                        .get_python_type()
                     )
                 else:
-                    program.import_(
-                        f"{target_pkg.module(prop.target.get_pymodule_name()).path}.{prop.target.name}",
-                        is_import_attr=True,
+                    pytype = PyTypeWithDep(
+                        f"Upsert{prop.target.name}",
+                        f"{target_pkg.module(prop.target.get_pymodule_name()).path}.Upsert{prop.target.name}",
                     )
-                    pytype = prop.target.name
+
+                if pytype.dep is not None:
+                    program.import_(pytype.dep, True)
 
                 if prop.cardinality.is_star_to_many():
-                    pytype = f"list[{pytype}]"
-                cls_ast(stmt.DefClassVarStatement(prop.name, pytype))
+                    pytype = pytype.as_list_type()
+                cls_ast(stmt.DefClassVarStatement(prop.name, pytype.type))
 
         # has_to_db = True
         # if any(prop for prop in cls.properties.values() if isinstance(prop, ObjectProperty) and prop.cardinality == Cardinality.MANY_TO_MANY):
@@ -128,11 +174,15 @@ def make_python_data_model(schema: Schema, target_pkg: Package):
                 [
                     DeferredVar.simple("self"),
                 ],
-                return_type=expr.ExprIdent(f"{cls.name}DB"),
+                return_type=expr.ExprIdent(
+                    f"{cls.name}DB" if cls.db is not None else cls.name
+                ),
             )(
                 lambda ast10: ast10.return_(
                     expr.ExprFuncCall(
-                        expr.ExprIdent(f"{cls.name}DB"),
+                        expr.ExprIdent(
+                            f"{cls.name}DB" if cls.db is not None else cls.name
+                        ),
                         [
                             to_db_type_conversion(
                                 program, expr.ExprIdent("self"), cls, prop
@@ -144,83 +194,120 @@ def make_python_data_model(schema: Schema, target_pkg: Package):
             ),
         )
 
-    def make_update(program: Program, cls: Class):
+    def make_normal(program: Program, cls: Class):
         if not cls.is_public:
             # skip classes that are not public
             return
 
         program.import_("__future__.annotations", True)
         program.import_("msgspec", False)
-        program.import_(
-            app.models.db.path + f".{cls.get_pymodule_name()}.{cls.name}",
-            True,
-            alias=f"{cls.name}DB",
-        )
+        if cls.db is not None:
+            # if the class is stored in the database, we need to import the database module
+            program.import_(
+                app.models.db.path + f".{cls.get_pymodule_name()}.{cls.name}",
+                True,
+                alias=f"{cls.name}DB",
+            )
 
         cls_ast = program.root.class_(cls.name, [expr.ExprIdent("msgspec.Struct")])
         for prop in cls.properties.values():
-            if prop.is_private:
+            if prop.data.is_private:
                 # skip private fields as this is for APIs exchange
                 continue
 
             if isinstance(prop, DataProperty):
-                pytype = prop.datatype.get_python_type()
+                pytype = prop.get_data_model_datatype().get_python_type()
                 if pytype.dep is not None:
                     program.import_(pytype.dep, True)
                 cls_ast(stmt.DefClassVarStatement(prop.name, pytype.type))
             elif isinstance(prop, ObjectProperty):
-                program.import_(
-                    f"{target_pkg.module(prop.target.get_pymodule_name()).path}.{prop.target.name}",
-                    is_import_attr=True,
-                )
-                if prop.cardinality.is_star_to_many():
-                    pytype = f"list[{prop.target.name}]"
+                if prop.target.db is not None:
+                    pytype = (
+                        assert_not_null(prop.target.get_id_property())
+                        .get_data_model_datatype()
+                        .get_python_type()
+                    )
                 else:
-                    pytype = prop.target.name
-                cls_ast(stmt.DefClassVarStatement(prop.name, pytype))
+                    pytype = PyTypeWithDep(
+                        prop.target.name,
+                        f"{target_pkg.module(prop.target.get_pymodule_name()).path}.{prop.target.name}",
+                    )
+
+                if pytype.dep is not None:
+                    program.import_(pytype.dep, True)
+
+                if prop.cardinality.is_star_to_many():
+                    pytype = pytype.as_list_type()
+                cls_ast(stmt.DefClassVarStatement(prop.name, pytype.type))
 
         cls_ast(
             stmt.LineBreak(),
-            stmt.PythonDecoratorStatement(
-                expr.ExprFuncCall(expr.ExprIdent("classmethod"), [])
+            (
+                stmt.PythonDecoratorStatement(
+                    expr.ExprFuncCall(expr.ExprIdent("classmethod"), [])
+                )
+                if cls.db is not None
+                else None
             ),
-            lambda ast00: ast00.func(
-                "from_db",
-                [
-                    DeferredVar.simple("cls"),
-                    DeferredVar.simple("record", expr.ExprIdent(f"{cls.name}DB")),
-                ],
-            )(
-                lambda ast10: ast10.return_(
-                    expr.ExprFuncCall(
-                        expr.ExprIdent("cls"),
-                        [
-                            from_db_type_conversion(expr.ExprIdent("record"), prop)
-                            for prop in cls.properties.values()
-                            if not prop.is_private
-                        ],
+            lambda ast00: (
+                ast00.func(
+                    "from_db",
+                    [
+                        DeferredVar.simple("cls"),
+                        DeferredVar.simple("record", expr.ExprIdent(f"{cls.name}DB")),
+                    ],
+                )(
+                    lambda ast10: ast10.return_(
+                        expr.ExprFuncCall(
+                            expr.ExprIdent("cls"),
+                            [
+                                from_db_type_conversion(expr.ExprIdent("record"), prop)
+                                for prop in cls.properties.values()
+                                if not prop.data.is_private
+                            ],
+                        )
                     )
                 )
+                if cls.db is not None
+                else None
             ),
         )
 
     for cls in schema.topological_sort():
+        if cls.name in reference_classes:
+            continue
+
         program = Program()
-        make_create(program, cls)
+        make_upsert(program, cls)
         program.root.linebreak()
-        make_update(program, cls)
+        make_normal(program, cls)
         target_pkg.module(cls.get_pymodule_name()).write(program)
 
 
 def make_python_relational_model(
-    schema: Schema, target_pkg: Package, target_data_pkg: Package
+    schema: Schema,
+    target_pkg: Package,
+    target_data_pkg: Package,
+    reference_classes: dict[str, ObjectPath],
 ):
     """Make python classes for relational database using SQLAlchemy.
 
     The new classes is going be compatible with SQLAlchemy 2.
+
+    Args:
+        schema: The schema to generate the classes from.
+        target_pkg: The package to write the classes to.
+        target_data_pkg: The package to write the data classes to.
+        reference_classes: A dictionary of class names to their references (e.g., the ones that are defined outside and used as referenced such as Tenant).
     """
     app = target_pkg.app
 
+    def get_property_name(prop: DataProperty | ObjectProperty):
+        if isinstance(prop, ObjectProperty):
+            if prop.target.db is not None:
+                return f"{prop.name}_id"
+        return prop.name
+
     def make_base(custom_types: Sequence[ObjectProperty]):
         """Make a base class for our database."""
         program = Program()
@@ -246,18 +333,18 @@ def make_python_relational_model(
             if custom_type.cardinality.is_star_to_many():
                 if custom_type.is_map:
                     program.import_("typing.Mapping", True)
-                    program.import_("sera.libs.baseorm.DictDataClassType", True)
+                    program.import_("sera.libs.base_orm.DictDataclassType", True)
                     type = f"Mapping[str, {custom_type.target.name}]"
-                    maptype = f"DictDataClassType({custom_type.target.name})"
+                    maptype = f"DictDataclassType({custom_type.target.name})"
                 else:
                     program.import_("typing.Sequence", True)
-                    program.import_("sera.libs.baseorm.ListDataClassType", True)
+                    program.import_("sera.libs.base_orm.ListDataclassType", True)
                     type = f"Sequence[str, {custom_type.target.name}]"
-                    maptype = f"ListDataClassType({custom_type.target.name})"
+                    maptype = f"ListDataclassType({custom_type.target.name})"
             else:
-                program.import_("sera.libs.baseorm.DataClassType", True)
+                program.import_("sera.libs.base_orm.DataclassType", True)
                 type = custom_type.target.name
-                maptype = f"DataClassType({custom_type.target.name})"
+                maptype = f"DataclassType({custom_type.target.name})"
 
             if custom_type.is_optional:
                 program.import_("typing.Optional", True)
@@ -305,7 +392,7 @@ def make_python_relational_model(
     custom_types: list[ObjectProperty] = []
 
     for cls in schema.topological_sort():
-        if cls.db is None:
+        if cls.db is None or cls.name in reference_classes:
             # skip classes that are not stored in the database
             continue
 
@@ -316,6 +403,40 @@ def make_python_relational_model(
         program.import_("sqlalchemy.orm.Mapped", True)
         program.import_(f"{target_pkg.path}.base.Base", True)
 
+        index_stmts = []
+        if len(cls.db.indices) > 0:
+            program.import_("sqlalchemy.Index", True)
+            index_stmts.append(
+                stmt.DefClassVarStatement(
+                    "_table_args__",
+                    None,
+                    PredefinedFn.tuple(
+                        [
+                            expr.ExprFuncCall(
+                                expr.ExprIdent("Index"),
+                                [expr.ExprConstant(index.name)]
+                                + [
+                                    expr.ExprConstant(
+                                        get_property_name(cls.properties[prop])
+                                    )
+                                    for prop in index.columns
+                                ]
+                                + (
+                                    [
+                                        PredefinedFn.keyword_assignment(
+                                            "unique", expr.ExprConstant(index.unique)
+                                        )
+                                    ]
+                                    if index.unique
+                                    else []
+                                ),
+                            )
+                            for index in cls.db.indices
+                        ]
+                    ),
+                )
+            )
+
         cls_ast = program.root.class_(
             cls.name, [expr.ExprIdent("MappedAsDataclass"), expr.ExprIdent("Base")]
         )
@@ -325,6 +446,7 @@ def make_python_relational_model(
                 type=None,
                 value=expr.ExprConstant(cls.db.table_name),
             ),
+            *index_stmts,
             stmt.LineBreak(),
         )
 
@@ -334,14 +456,14 @@ def make_python_relational_model(
                 continue
 
             if isinstance(prop, DataProperty):
-                pytype = prop.datatype.get_sqlalchemy_type()
-                if pytype.dep is not None:
-                    program.import_(pytype.dep, True)
+                sqltype = prop.datatype.get_sqlalchemy_type()
+                for dep in sqltype.deps:
+                    program.import_(dep, True)
 
                 propname = prop.name
-                proptype = f"Mapped[{pytype.type}]"
+                proptype = f"Mapped[{sqltype.mapped_pytype}]"
 
-                propvalargs = []
+                propvalargs: list[expr.Expr] = [expr.ExprIdent(sqltype.type)]
                 if prop.db.is_primary_key:
                     propvalargs.append(
                         PredefinedFn.keyword_assignment(
@@ -354,12 +476,25 @@ def make_python_relational_model(
                                 "autoincrement", expr.ExprConstant("auto")
                             )
                         )
+                else:
                     if prop.db.is_unique:
                         propvalargs.append(
                             PredefinedFn.keyword_assignment(
                                 "unique", expr.ExprConstant(True)
                             )
                         )
+                    elif prop.db.is_indexed:
+                        propvalargs.append(
+                            PredefinedFn.keyword_assignment(
+                                "index", expr.ExprConstant(True)
+                            )
+                        )
+                if prop.db.is_nullable:
+                    propvalargs.append(
+                        PredefinedFn.keyword_assignment(
+                            "nullable", expr.ExprConstant(True)
+                        )
+                    )
                 propval = expr.ExprFuncCall(
                     expr.ExprIdent("mapped_column"), propvalargs
                 )
@@ -369,6 +504,7 @@ def make_python_relational_model(
                 make_python_relational_object_property(
                     program=program,
                     target_pkg=target_pkg,
+                    target_data_pkg=target_data_pkg,
                     cls_ast=cls_ast,
                     cls=cls,
                     prop=prop,
@@ -387,6 +523,7 @@ def make_python_relational_model(
 def make_python_relational_object_property(
     program: Program,
     target_pkg: Package,
+    target_data_pkg: Package,
     cls_ast: AST,
     cls: Class,
     prop: ObjectProperty,
@@ -395,7 +532,7 @@ def make_python_relational_object_property(
     assert prop.db is not None
     if prop.target.db is not None:
         # if the target class is in the database, we generate a foreign key for it.
-        program.import_("sqlalchemy.orm.relationship", True)
+        program.import_("sqlalchemy.ForeignKey", True)
 
         if prop.cardinality == Cardinality.MANY_TO_MANY:
             make_python_relational_object_property_many_to_many(
@@ -411,13 +548,14 @@ def make_python_relational_object_property(
         idprop = prop.target.get_id_property()
         assert idprop is not None
         idprop_pytype = idprop.datatype.get_sqlalchemy_type()
-        if idprop_pytype.dep is not None:
-            program.import_(idprop_pytype.dep, True)
+        for dep in idprop_pytype.deps:
+            program.import_(dep, True)
 
-        proptype = f"Mapped[{idprop_pytype.type}]"
+        proptype = f"Mapped[{idprop_pytype.mapped_pytype}]"
         propval = expr.ExprFuncCall(
             expr.ExprIdent("mapped_column"),
             [
+                expr.ExprIdent(idprop_pytype.type),
                 expr.ExprFuncCall(
                     expr.ExprIdent("ForeignKey"),
                     [
@@ -440,7 +578,7 @@ def make_python_relational_object_property(
 
     # if the target class is not in the database,
     program.import_(
-        f"{target_pkg.module(prop.target.get_pymodule_name()).path}.{prop.target.name}",
+        f"{target_data_pkg.module(prop.target.get_pymodule_name()).path}.{prop.target.name}",
         is_import_attr=True,
     )
     propname = prop.name
@@ -474,6 +612,7 @@ def make_python_relational_object_property(
             propvalargs,
         )
     else:
+        assert prop.db.is_embedded == "json"
         # we create a custom field, the custom field mapping need to be defined in the base
         propval = expr.ExprFuncCall(expr.ExprIdent("mapped_column"), [])
         custom_types.append(prop)
@@ -527,21 +666,6 @@ def make_python_relational_object_property_many_to_many(
                 value=expr.ExprConstant(f"{clsdb.table_name}_{targetdb.table_name}"),
             ),
             stmt.LineBreak(),
-            stmt.DefClassVarStatement(
-                "id",
-                "Mapped[int]",
-                expr.ExprFuncCall(
-                    expr.ExprIdent("mapped_column"),
-                    [
-                        PredefinedFn.keyword_assignment(
-                            "primary_key", expr.ExprConstant(True)
-                        ),
-                        PredefinedFn.keyword_assignment(
-                            "autoincrement", expr.ExprConstant("auto")
-                        ),
-                    ],
-                ),
-            ),
             stmt.DefClassVarStatement(
                 to_snake_case(cls.name),
                 f"Mapped[{cls.name}]",
@@ -585,6 +709,9 @@ def make_python_relational_object_property_many_to_many(
                                 ),
                             ],
                         ),
+                        PredefinedFn.keyword_assignment(
+                            "primary_key", expr.ExprConstant(True)
+                        ),
                     ],
                 ),
             ),
@@ -627,6 +754,9 @@ def make_python_relational_object_property_many_to_many(
                                 ),
                             ],
                         ),
+                        PredefinedFn.keyword_assignment(
+                            "primary_key", expr.ExprConstant(True)
+                        ),
                     ],
                 ),
             ),
@@ -639,6 +769,8 @@ def make_python_relational_object_property_many_to_many(
     # now we add the relationship to the source.
     # we can configure it to be list, set, or dict depends on what we want.
     program.import_(new_table_module.path + f".{new_table}", True)
+    program.import_("sqlalchemy.orm.relationship", True)
+
     # program.import_("typing.TYPE_CHECKING", True)
     # program.import_area.if_(expr.ExprIdent("TYPE_CHECKING"))(
     #     lambda ast00: ast00.import_(
@@ -666,3 +798,13 @@ def make_python_relational_object_property_many_to_many(
             ),
         ),
     )
+
+
+def get_data_conversion(
+    source_pytype: str, target_pytype: str
+) -> Callable[[expr.Expr], expr.Expr]:
+    if source_pytype == target_pytype:
+        return lambda x: x
+    if source_pytype == "str" and target_pytype == "bytes":
+        return lambda x: expr.ExprMethodCall(x, "encode", [])
+    raise NotImplementedError(f"Cannot convert {source_pytype} to {target_pytype}")