➕add sqlmodel export flag and primary key

custom_templates/BaseModel.jinja2

@@ -1,7 +1,7 @@
 {% for decorator in decorators -%}
 {{ decorator }}
 {% endfor -%}
-class {{ class_name }}({% if base_class %}{{ base_class }}{% endif %}, table=True):{% if comment is defined %}  # {{ comment }}{% endif %}
+class {{ class_name }}({% if base_class %}{{ base_class }}{% endif %}{% if SQL %}, table=True{% endif %}):{% if comment is defined %}  # {{ comment }}{% endif %}
 {%- if description %}
     """
     {{ description | indent(4) }}
@@ -37,3 +37,5 @@ class {{ class_name }}({% if base_class %}{{ base_class }}{% endif %}, table=Tru
     {{ method }}
 {%- endfor -%}
 {%- endfor -%}
+{%- if SQL and SQL['primary'] %}
+    {{ SQL['primary'] }}{%- endif %}

requirements.txt

@@ -25,7 +25,7 @@ click==8.1.7
 colorama==0.4.6
     # via click
 datamodel-code-generator==0.23.0
-    # via -r .\requirements.in
+    # via -r requirements.in
 dnspython==2.4.2
     # via email-validator
 email-validator==2.0.0.post2

src/bo4e_generator/__main__.py

@@ -20,13 +20,15 @@ def resolve_paths(input_directory: Path, output_directory: Path) -> tuple[Path,
     return input_directory, output_directory
 
 
-def generate_bo4e_schemas(input_directory: Path, output_directory: Path, pydantic_v1: bool = False):
+def generate_bo4e_schemas(
+    input_directory: Path, output_directory: Path, pydantic_v1: bool = False, sql_model: bool = False
+):
     """
     Generate all BO4E schemas from the given input directory and save them in the given output directory.
     """
     input_directory, output_directory = resolve_paths(input_directory, output_directory)
     namespace = get_namespace(input_directory)
-    file_contents = parse_bo4e_schemas(input_directory, namespace, pydantic_v1)
+    file_contents = parse_bo4e_schemas(input_directory, namespace, pydantic_v1, sql_model)
     file_contents[Path("__init__.py")] = bo4e_init_file_content(get_version(namespace))
     for relative_file_path, file_content in file_contents.items():
         file_path = output_directory / relative_file_path
@@ -60,13 +62,21 @@ def generate_bo4e_schemas(input_directory: Path, output_directory: Path, pydanti
     required=False,
     default=False,
 )
+@click.option(
+    "--sql-model",
+    "-sqlm",
+    is_flag=True,
+    help="Generate SQLModel classes.",
+    required=False,
+    default=False,
+)
 @click.help_option()
 @click.version_option(package_name="BO4E-Python-Generator")
-def main(input_dir: Path, output_dir: Path, pydantic_v1: bool):
+def main(input_dir: Path, output_dir: Path, pydantic_v1: bool, sql_model: bool):
     """
     CLI entry point for the bo4e-generator.
     """
-    generate_bo4e_schemas(input_dir, output_dir, pydantic_v1)
+    generate_bo4e_schemas(input_dir, output_dir, pydantic_v1, sql_model)
 
 
 if __name__ == "__main__":

src/bo4e_generator/parser.py

@@ -2,7 +2,9 @@
 Contains code to generate pydantic v2 models from json schemas.
 Since the used tool doesn't support all features we need, we monkey patch some functions.
 """
+import json
 import re
+from collections import defaultdict
 from pathlib import Path
 from typing import Tuple
 
@@ -119,7 +121,7 @@ def remove_future_import(python_code: str) -> str:
 
 
 def parse_bo4e_schemas(
-    input_directory: Path, namespace: dict[str, SchemaMetadata], pydantic_v1: bool = False
+    input_directory: Path, namespace: dict[str, SchemaMetadata], pydantic_v1: bool = False, sql_model: bool = False
 ) -> dict[Path, str]:
     """
     Generate all BO4E schemas from the given input directory. Returns all file contents as dictionary:
@@ -132,9 +134,23 @@ def parse_bo4e_schemas(
     )
     monkey_patch_relative_import()
 
+    additional_arguments = {}
+
+    if sql_model:
+        additional_sql_data = defaultdict(dict)
+        for schema_metadata in namespace.values():
+            if schema_metadata.pkg != "enum":
+                additional_sql_data[schema_metadata.class_name]["SQL"] = {
+                    "primary": schema_metadata.class_name.lower()
+                    + "_id: Field( default_factory=uuid_pkg.uuid4, primary_key=True, index=True, nullable=False )"
+                }
+        additional_arguments["extra_template_data"] = additional_sql_data
+        additional_arguments["additional_imports"] = ["sqlmodel.Field"]
+        additional_arguments["base_class"] = "sqlmodel.SQLModel"
+        additional_arguments["custom_template_dir"] = Path.cwd() / Path(".\custom_templates")
+
     parser = JsonSchemaParser(
         input_directory,
-        base_class="sqlmodel.SQLModel",
         data_model_type=data_model_types.data_model,
         data_model_root_type=data_model_types.root_model,
         data_model_field_type=data_model_types.field_model,
@@ -153,7 +169,7 @@ def parse_bo4e_schemas(
         base_path=input_directory,
         remove_special_field_name_prefix=True,
         allow_extra_fields=False,
-        custom_template_dir=Path.cwd() / Path(".\custom_templates"),
+        **additional_arguments,
     )
     parse_result = parser.parse()
     if not isinstance(parse_result, dict):

src/bo4e_generator/schema.py

@@ -34,6 +34,20 @@ def save(self, content: str):
         self.output_file.parent.mkdir(parents=True, exist_ok=True)
         self.output_file.write_text(content)
 
+    def add_uuid(self):
+        """
+        add UUID field classname_id for usage with sqlmodel
+        """
+        self.schema_parsed["properties"][self.class_name.lower() + "_id"] = {
+            "pydantic_field": {
+                "default_factory": "uuid_pkg.uuid4",
+                "primary_key": True,
+                "index": True,
+                "nullable": False,
+            },
+            "type": "uuid.UUID",
+        }
+
     def __str__(self):
         return f"{self.pkg}.{self.class_name}"