Merge pull request #14 from sebcampos/issue13fix

fix for _get_schema_v2 always returning None
eadwinCode · Nov 18, 2024 · 150d36a · 150d36a
2 parents 0df47c8 + e2a51c9
commit 150d36a
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 5 deletions.
diff --git a/ninja_schema/orm/factory.py b/ninja_schema/orm/factory.py
@@ -89,10 +89,10 @@ def _get_schema_v2(
         cls, name: str, model_config_kwargs: typing.Dict, model_type: typing.Type
     ) -> Union[Type["ModelSchema"], Type["Schema"]]:
         model_config = cls.get_model_config(**model_config_kwargs)
-
+        new_schema_result = {}
         new_schema_string = f"""class {name}(model_type):
             class Config(model_config):
                 pass """
 
-        exec(new_schema_string, locals())
-        return locals().get(name)  # type:ignore[return-value]
+        exec(new_schema_string, locals(), new_schema_result)
+        return new_schema_result.get(name)  # type:ignore[return-value]
diff --git a/ninja_schema/pydanticutils/__init__.py b/ninja_schema/pydanticutils/__init__.py
@@ -15,6 +15,7 @@
 
 PYDANTIC_VERSION = list(map(int, _PYDANTIC_VERSION.split(".")))[:2]
 IS_PYDANTIC_V1 = PYDANTIC_VERSION[0] == 1
+IS_PYDANTIC_V292_OR_GREATER = PYDANTIC_VERSION[0] >= 2 and PYDANTIC_VERSION[1] >= 9
 
 
 def is_valid_field_name(name: str) -> bool:

diff --git a/tests/test_v2_pydantic/test_custom_fields.py b/tests/test_v2_pydantic/test_custom_fields.py
@@ -4,12 +4,12 @@
 from pydantic import ValidationError
 
 from ninja_schema import ModelSchema
-from ninja_schema.pydanticutils import IS_PYDANTIC_V1
+from ninja_schema.pydanticutils import IS_PYDANTIC_V1, IS_PYDANTIC_V292_OR_GREATER
 from tests.models import Student, StudentEmail
 
 
 class TestCustomFields:
-    @pytest.mark.skipif(IS_PYDANTIC_V1, reason="requires pydantic == 2.1.x")
+    @pytest.mark.skipif(IS_PYDANTIC_V1 and not IS_PYDANTIC_V292_OR_GREATER, reason="requires pydantic == 2.1.x and pydantic < 2.9.2")
     def test_enum_field(self):
         class StudentSchema(ModelSchema):
             model_config = {"model": Student, "include": "__all__"}
@@ -45,6 +45,44 @@ class StudentSchema(ModelSchema):
         with pytest.raises(ValidationError):
             StudentSchema(semester="something")
 
+
+    @pytest.mark.skipif(not IS_PYDANTIC_V292_OR_GREATER, reason="requires pydantic == 2.1.x and pydantic < 2.9.2")
+    def test_enum_field_v292_or_greater(self):
+        class StudentSchema(ModelSchema):
+            model_config = {"model": Student, "include": "__all__"}
+
+        print(json.dumps(StudentSchema.schema(), sort_keys=False, indent=4))
+        assert StudentSchema.schema() == {
+            "$defs": {
+                "SemesterEnum": {
+                    "enum": ["1", "2", "3"],
+                    "title": "SemesterEnum",
+                    "type": "string",
+                }
+            },
+            "properties": {
+                "id": {
+                    "anyOf": [{"type": "integer"}, {"type": "null"}],
+                    "default": None,
+                    "description": "",
+                    "title": "Id",
+                },
+                "semester": {
+                    "$ref": "#/$defs/SemesterEnum",
+                    "default": "1",
+                    "description": "",
+                    "title": "Semester",
+                },
+            },
+            "title": "StudentSchema",
+            "type": "object",
+        }
+        schema_instance = StudentSchema(semester="1")
+        assert str(schema_instance.json()) == '{"id":null,"semester":"1"}'
+        with pytest.raises(ValidationError):
+            StudentSchema(semester="something")
+
+
     @pytest.mark.skipif(IS_PYDANTIC_V1, reason="requires pydantic == 2.1.x")
     def test_email_field(self):
         class StudentEmailSchema(ModelSchema):