CCM-12896: Generate Pydantic models programmatically

Author: gareth-allan

src/python-schema-generator/src/file_utils.py

@@ -24,17 +24,29 @@ def list_json_schemas(schema_dir: str) -> list[str]:
     return [f.name for f in flattened_schema_files]
 
 
-def load_json_schema(schema_path: str) -> dict[str, Any]:
-    """Load a JSON schema from file.
+def load_json_schema(schema_path: str) -> str:
+    """Load a JSON schema from file as a string.
 
     Args:
         schema_path: Path to the JSON schema file
 
     Returns:
-        Loaded schema as dictionary
+        Loaded schema as string
     """
     with open(schema_path, encoding="utf-8") as f:
-        return json.load(f)
+        return f.read()
+
+
+def parse_json_schema(schema: str) -> dict[str, Any]:
+    """Parse a JSON schema from a string.
+
+    Args:
+        schema: JSON schema as a string
+
+    Returns:
+        Schema as dictionary
+    """
+    return json.loads(schema)
 
 
 def write_init_file(output_dir: str, model_names: list[str]) -> None:

src/python-schema-generator/src/generate_models.py

@@ -13,6 +13,7 @@
 from file_utils import (
     list_json_schemas,
     load_json_schema,
+    parse_json_schema,
     model_name_to_module_name,
     write_init_file,
 )
@@ -62,14 +63,15 @@ def main() -> int:
         generated_models = []
         for schema_filename in schema_filenames:
             schema_path = str(Path(args.input_dir) / schema_filename)
-            schema = load_json_schema(schema_path)
+            string_schema = load_json_schema(schema_path)
+            schema = parse_json_schema(string_schema)
 
             model_name = extract_model_name(schema)
             output_filename = model_name_to_module_name(model_name) + ".py"
-            output_file_path = str(Path(args.output_dir) / output_filename)
+            output_file_path = Path(args.output_dir) / output_filename
 
             generate_pydantic_model(
-                schema_path, output_file_path, model_name
+                string_schema, output_file_path, model_name
             )
 
             generated_models.append(model_name)

src/python-schema-generator/src/model_generator.py

@@ -1,13 +1,12 @@
 """Model generator using datamodel-code-generator."""
 
-import subprocess
-import sys
 from pathlib import Path
+from datamodel_code_generator import DataModelType, InputFileType, generate
 
 
 
 def generate_pydantic_model(
-    schema_path: str, output_file: str, class_name: str
+    schema: str, output_file_path: Path, class_name: str
 ) -> None:
     """Generate a Pydantic model from a JSON schema.
 
@@ -19,32 +18,18 @@ def generate_pydantic_model(
     Raises:
         RuntimeError: If model generation fails
     """
-    datamodel_cmd = str(Path(sys.executable).parent / "datamodel-codegen")
-    cmd = [
-        datamodel_cmd,
-        "--input",
-        schema_path,
-        "--output",
-        output_file,
-        "--class-name",
-        class_name,
-        "--input-file-type",
-        "jsonschema",
-        "--output-model-type",
-        "pydantic_v2.BaseModel",
-        "--use-schema-description",
-        "--custom-file-header",
-        '''"""Generated Pydantic model for NHS Notify Digital Letters events.
+
+    generate(
+        schema,
+        input_file_type=InputFileType.JsonSchema,
+        output=output_file_path,
+        output_model_type=DataModelType.PydanticV2BaseModel,
+        class_name=class_name,
+        use_schema_description=True,
+        custom_file_header='''"""Generated Pydantic model for NHS Notify Digital Letters events.
 
 This file is auto-generated. Do not edit manually.
 """
 
 '''
-    ]
-    result = subprocess.run(
-        cmd, capture_output=True, text=True, check=False, encoding="utf-8"
     )
-
-    if result.returncode != 0:
-        error_msg = f"Failed to generate model: {result.stderr}"
-        raise RuntimeError(error_msg)

src/python-schema-generator/tests/test_file_utils.py

@@ -7,6 +7,7 @@
 from src.file_utils import (
     list_json_schemas,
     load_json_schema,
+    parse_json_schema,
     write_init_file,
     model_name_to_module_name
 )
@@ -51,21 +52,31 @@ class TestLoadJsonSchema:
     def test_loads_valid_json_schema(self, tmp_path):
         """Test loading a valid JSON schema."""
         schema_file = tmp_path / "test.schema.json"
-        schema_content = {"title": "TestSchema", "type": "object"}
-        schema_file.write_text(json.dumps(schema_content))
+        schema_content = json.dumps({"title": "TestSchema", "type": "object"})
+        schema_file.write_text(schema_content)
 
         result = load_json_schema(str(schema_file))
 
         assert result == schema_content
 
-    def test_raises_error_for_invalid_json(self, tmp_path):
+
+class TestParseJsonSchema:
+    """Tests for parse_json_schema function."""
+
+    def test_parses_valid_json_schema(self):
+        """Test loading a valid JSON schema."""
+        schema_content = { "title": "TestSchema", "type": "object" }
+
+        result = parse_json_schema(json.dumps(schema_content))
+
+        assert result == schema_content
+
+    def test_raises_error_for_invalid_json(self):
         """Test that it raises error for invalid JSON."""
-        schema_file = tmp_path / "invalid.json"
-        schema_file.write_text("not valid json")
+        invalid_json = "not valid json"
 
         with pytest.raises(json.JSONDecodeError):
-            load_json_schema(str(schema_file))
-
+            parse_json_schema(invalid_json)
 
 class TestWriteInitFile:
     """Tests for write_init_file function."""

src/python-schema-generator/tests/test_model_generator.py

@@ -1,45 +1,49 @@
 """Tests for model_generator module."""
 
-from unittest.mock import MagicMock, patch
+from pathlib import Path
+from unittest.mock import patch
 
 import pytest
+from datamodel_code_generator import DataModelType, InputFileType
 
 from src.model_generator import generate_pydantic_model
 
 
 class TestGeneratePydanticModel:
     """Tests for generate_pydantic_model function."""
 
-    @patch("src.model_generator.subprocess.run")
-    def test_calls_datamodel_codegen_with_expected_arguments(self, mock_run):
+    @patch("src.model_generator.generate")
+    def test_calls_datamodel_codegen_with_expected_arguments(self, mock_generate):
         """Test successful model generation."""
         # Arrange
-        schema_path = "test_model.json"
-        output_file = "test_model.py"
-        mock_run.return_value = MagicMock(returncode=0, stderr="")
+        schema = '{"type": "object", "properties": {"name": {"type": "string"}}}'
+        output_file = Path("test_model.py")
 
         # Act
-        generate_pydantic_model(schema_path, output_file, "TestModel")
+        generate_pydantic_model(schema, output_file, "TestModel")
 
         # Assert
-        assert mock_run.called
-        cmd_args = mock_run.call_args[0][0]
-        assert "datamodel-codegen" in cmd_args[0]  # First arg is the executable
-        assert "--input" in cmd_args[1]
-        assert schema_path in cmd_args[2]
-        assert "--output" in cmd_args[3]
-        assert output_file in cmd_args[4]
-        assert "--class-name" in cmd_args[5]
-        assert "TestModel" in cmd_args[6]
-
-    @patch("src.model_generator.subprocess.run")
-    def test_raises_error_on_generation_failure(self, mock_run):
-        """Test that it raises error when generation fails."""
-        schema_path = "test_model.json"
-        output_file = "test_model.py"
-        mock_run.return_value = MagicMock(
-            returncode=1, stderr="Error: Invalid schema"
+        mock_generate.assert_called_once_with(
+            schema,
+            input_file_type=InputFileType.JsonSchema,
+            output=output_file,
+            output_model_type=DataModelType.PydanticV2BaseModel,
+            class_name="TestModel",
+            use_schema_description=True,
+            custom_file_header='''"""Generated Pydantic model for NHS Notify Digital Letters events.
+
+This file is auto-generated. Do not edit manually.
+"""
+
+'''
         )
 
-        with pytest.raises(RuntimeError, match="Failed to generate model"):
-            generate_pydantic_model(schema_path, output_file, "TestModel")
+    @patch("src.model_generator.generate")
+    def test_raises_error_on_generation_failure(self, mock_generate):
+        """Test that it raises error when generation fails."""
+        schema = '{"type": "object", "properties": {"name": {"type": "string"}}}'
+        output_file = Path("test_model.py")
+        mock_generate.side_effect = RuntimeError("Invalid schema")
+
+        with pytest.raises(RuntimeError, match="Invalid schema"):
+            generate_pydantic_model(schema, output_file, "TestModel")