Improve Code to fit better with datamodel-code-generator (#13)

* Improve Code

* 🩹

* 🩹

src/bo4e_generator/__main__.py

@@ -5,7 +5,7 @@
 
 import click
 
-from bo4e_generator.parser import create_init_files, generate_bo4e_schema
+from bo4e_generator.parser import bo4e_init_file_content, parse_bo4e_schemas
 from bo4e_generator.schema import get_namespace, get_version
 
 
@@ -25,13 +25,16 @@ def generate_bo4e_schemas(input_directory: Path, output_directory: Path, pydanti
     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, output_directory)
-    for schema_metadata in namespace.values():
-        result = generate_bo4e_schema(schema_metadata, namespace, pydantic_v1)
-        schema_metadata.save(result)
-        print(f"Generated {schema_metadata}")
-    create_init_files(output_directory, get_version(namespace))
-    print(f"Generated __init__.py files in {output_directory}")
+    namespace = get_namespace(input_directory)
+    file_contents = parse_bo4e_schemas(input_directory, namespace, pydantic_v1)
+    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
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+        file_path.write_text(file_content)
+        print(f"Created {file_path}")
+
+    print("Done.")
 
 
 @click.command()

src/bo4e_generator/parser.py

@@ -36,13 +36,6 @@ def get_bo4e_data_model_types(
     def _module_path(self) -> list[str]:
         if self.name not in namespace:
             raise ValueError(f"Model not in namespace: {self.name}")
-        if self.file_path:
-            return [
-                *self.file_path.parts[:-1],
-                self.file_path.stem,
-                namespace[self.name].pkg,
-                namespace[self.name].module_name,
-            ]
         return [namespace[self.name].pkg, namespace[self.name].module_name]
 
     @property  # type: ignore[misc]
@@ -99,7 +92,7 @@ def _get_valid_name(
 
         name = re.sub(r"[¹²³⁴⁵⁶⁷⁸⁹]|\W", "_", name)
         if name[0].isnumeric():
-            name = f"{self.special_field_name_prefix}{name}"  # Changed line by me
+            name = f"{self.special_field_name_prefix}_{name}"  # Changed line by me
 
         # We should avoid having a field begin with an underscore, as it
         # causes pydantic to consider it as private
@@ -175,15 +168,11 @@ def relative(current_module: str, reference: str) -> Tuple[str, str]:
     datamodel_code_generator.parser.base.relative = relative
 
 
-def create_init_files(output_path: Path, version: str) -> None:
+def bo4e_init_file_content(version: str) -> str:
     """
     Create __init__.py files in all subdirectories of the given output directory and in the directory itself.
     """
-    (output_path / "__init__.py").write_text(
-        f'""" Contains information about the bo4e version """\n\n__version__ = "{version}"\n'
-    )
-    for directory in output_path.glob("**/"):
-        (directory / "__init__.py").touch()
+    return f'""" Contains information about the bo4e version """\n\n__version__ = "{version}"\n'
 
 
 def remove_future_import(python_code: str) -> str:
@@ -193,11 +182,12 @@ def remove_future_import(python_code: str) -> str:
     return re.sub(r"from __future__ import annotations\n\n", "", python_code)
 
 
-def generate_bo4e_schema(
-    schema_metadata: SchemaMetadata, namespace: dict[str, SchemaMetadata], pydantic_v1: bool = False
-) -> str:
+def parse_bo4e_schemas(
+    input_directory: Path, namespace: dict[str, SchemaMetadata], pydantic_v1: bool = False
+) -> dict[Path, str]:
     """
-    Generate a pydantic v2 model from the given schema. Returns the resulting code as string.
+    Generate all BO4E schemas from the given input directory. Returns all file contents as dictionary:
+    file path (relative to arbitrary output directory) => file content.
     """
     data_model_types = get_bo4e_data_model_types(
         DataModelType.PydanticBaseModel if pydantic_v1 else DataModelType.PydanticV2BaseModel,
@@ -208,7 +198,7 @@ def generate_bo4e_schema(
     monkey_patch_relative_import()
 
     parser = JsonSchemaParser(
-        schema_metadata.schema_text,
+        input_directory,
         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,
@@ -223,29 +213,32 @@ def generate_bo4e_schema(
         set_default_enum_member=True,
         snake_case_field=True,
         field_constraints=True,
-        class_name=schema_metadata.class_name,
         capitalise_enum_members=True,
-        base_path=schema_metadata.input_file.parent,
+        base_path=input_directory,
         remove_special_field_name_prefix=True,
-        special_field_name_prefix="field_",
         allow_extra_fields=False,
     )
-    result = parser.parse()
-    if isinstance(result, dict):
+    parse_result = parser.parse()
+    if not isinstance(parse_result, dict):
+        raise ValueError(f"Unexpected type of parse result: {type(parse_result)}")
+    file_contents = {}
+    for schema_metadata in namespace.values():
         if schema_metadata.module_name.startswith("_"):
             # Because somehow the generator uses the prefix also on the module name. Don't know why.
             module_path = (schema_metadata.pkg, f"field{schema_metadata.module_name}.py")
         else:
             module_path = (schema_metadata.pkg, f"{schema_metadata.module_name}.py")
-        try:
-            result = result[module_path].body
-        except KeyError as error:
+
+        if module_path not in parse_result:
             raise KeyError(
                 f"Could not find module {'.'.join(module_path)} in results: "
-                f"{list(result.keys())}"  # type: ignore[union-attr]
+                f"{list(parse_result.keys())}"  # type: ignore[union-attr]
                 # Item "str" of "str | dict[tuple[str, ...], Result]" has no attribute "keys"
                 # Somehow, mypy is not good enough to understand the instance-check above
-            ) from error
+            )
+
+        file_contents[schema_metadata.output_file] = remove_future_import(parse_result.pop(module_path).body)
+
+    file_contents.update({Path(*module_path): result.body for module_path, result in parse_result.items()})
 
-    result = remove_future_import(result)
-    return result
+    return file_contents

src/bo4e_generator/schema.py

@@ -19,10 +19,11 @@ class SchemaMetadata(BaseModel):
     schema_text: str
     schema_parsed: SchemaType
     class_name: str
-    input_file: Path
     pkg: str
     "e.g. 'bo'"
+    input_file: Path
     output_file: Path
+    "The output file will be a relative path"
     module_name: str
     "e.g. 'preisblatt_netznutzung"
 
@@ -45,7 +46,7 @@ def camel_to_snake(name: str) -> str:
     return re.sub("([a-z0-9])([A-Z])", r"\1_\2", name).lower()
 
 
-def get_namespace(input_directory: Path, output_directory: Path) -> dict[str, SchemaMetadata]:
+def get_namespace(input_directory: Path) -> dict[str, SchemaMetadata]:
     """
     Create a namespace for the bo4e classes.
     """
@@ -61,7 +62,7 @@ def get_namespace(input_directory: Path, output_directory: Path) -> dict[str, Sc
             pkg=file_path.parent.name,
             module_name=module_name,
             input_file=file_path,
-            output_file=output_directory / file_path.relative_to(input_directory).with_name(f"{module_name}.py"),
+            output_file=file_path.relative_to(input_directory).with_name(f"{module_name}.py"),
             schema_text=schema_text,
             schema_parsed=schema_parsed,
             class_name=class_name,

unittests/test_main.py

@@ -6,8 +6,6 @@
 from pydantic import BaseModel
 
 from bo4e_generator.__main__ import main
-from bo4e_generator.parser import generate_bo4e_schema
-from bo4e_generator.schema import get_namespace
 
 BASE_DIR = Path(__file__).parents[1]
 OUTPUT_DIR = Path("unittests/output/bo4e")
@@ -37,12 +35,3 @@ def test_main(self):
         from .output.bo4e import __version__  # type: ignore[import-not-found]
 
         assert __version__ == "0.6.1rc13"
-
-    def test_single(self):
-        os.chdir(BASE_DIR)
-        namespace = get_namespace(INPUT_DIR.resolve(), OUTPUT_DIR.resolve())
-        schema_metadata = namespace["Angebot"]
-        result = generate_bo4e_schema(schema_metadata, namespace)
-
-        assert "class Angebot(BaseModel):" in result
-        assert "    typ: Annotated[Typ" in result