From 5a4cbe78d1d2d8b09ed43a2dac58bf3bd68a2af4 Mon Sep 17 00:00:00 2001
From: Vitaliy Kucheryaviy <ppr.vitaly@gmail.com>
Date: Fri, 6 Dec 2024 19:00:10 +0200
Subject: [PATCH] add register_type public function (#1352)

* register_field utility

* mypy lint
---
 ninja/conf.py             |  2 +-
 ninja/orm/__init__.py     |  3 ++-
 ninja/orm/fields.py       | 15 ++++++++++++++-
 tests/test_orm_schemas.py | 21 ++++++++++++++++++++-
 4 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/ninja/conf.py b/ninja/conf.py
index 992322aef..1ae570e2d 100644
--- a/ninja/conf.py
+++ b/ninja/conf.py
@@ -11,7 +11,7 @@ class Settings(BaseModel):
         "ninja.pagination.LimitOffsetPagination", alias="NINJA_PAGINATION_CLASS"
     )
     PAGINATION_PER_PAGE: int = Field(100, alias="NINJA_PAGINATION_PER_PAGE")
-    PAGINATION_MAX_LIMIT: int = Field(inf, alias="NINJA_PAGINATION_MAX_LIMIT")
+    PAGINATION_MAX_LIMIT: int = Field(inf, alias="NINJA_PAGINATION_MAX_LIMIT")  # type: ignore
 
     # Throttling
     NUM_PROXIES: Optional[int] = Field(None, alias="NINJA_NUM_PROXIES")
diff --git a/ninja/orm/__init__.py b/ninja/orm/__init__.py
index 3924c9c6d..f75300d84 100644
--- a/ninja/orm/__init__.py
+++ b/ninja/orm/__init__.py
@@ -1,4 +1,5 @@
 from ninja.orm.factory import create_schema
+from ninja.orm.fields import register_field
 from ninja.orm.metaclass import ModelSchema
 
-__all__ = ["create_schema", "ModelSchema"]
+__all__ = ["create_schema", "register_field", "ModelSchema"]
diff --git a/ninja/orm/fields.py b/ninja/orm/fields.py
index 58ad478ec..d67814c8c 100644
--- a/ninja/orm/fields.py
+++ b/ninja/orm/fields.py
@@ -9,6 +9,7 @@
 from pydantic.fields import FieldInfo
 from pydantic_core import PydanticUndefined, core_schema
 
+from ninja.errors import ConfigError
 from ninja.openapi.schema import OpenAPISchema
 from ninja.types import DictStrAny
 
@@ -81,6 +82,10 @@ def validate(cls, value: Any, _: Any) -> Any:
 TModel = TypeVar("TModel")
 
 
+def register_field(django_field: str, python_type: Any) -> None:
+    TYPES[django_field] = python_type
+
+
 @no_type_check
 def create_m2m_link_type(type_: Type[TModel]) -> Type[TModel]:
     class M2MLink(type_):  # type: ignore
@@ -148,7 +153,15 @@ def get_schema_field(
         max_length = field_options.get("max_length")
 
         internal_type = field.get_internal_type()
-        python_type = TYPES[internal_type]
+        try:
+            python_type = TYPES[internal_type]
+        except KeyError as e:
+            msg = [
+                f"Do not know how to convert django field '{internal_type}'.",
+                "Try from ninja.orm import register_field",
+                f"register_field('{internal_type}', <your-python-type>)",
+            ]
+            raise ConfigError("\n".join(msg)) from e
 
         if field.primary_key or blank or null or optional:
             default = None
diff --git a/tests/test_orm_schemas.py b/tests/test_orm_schemas.py
index d3fde8c1f..5d8a05b48 100644
--- a/tests/test_orm_schemas.py
+++ b/tests/test_orm_schemas.py
@@ -8,7 +8,7 @@
 from util import pydantic_ref_fix
 
 from ninja.errors import ConfigError
-from ninja.orm import create_schema
+from ninja.orm import create_schema, register_field
 from ninja.orm.shortcuts import L, S
 
 
@@ -581,3 +581,22 @@ class Meta:
         SomeReqFieldModel, optional_fields=["some_field", "other_field", "optional"]
     )
     assert Schema.json_schema().get("required") is None
+
+
+def test_register_custom_field():
+    class MyCustomField(models.Field):
+        description = "MyCustomField"
+
+    class ModelWithCustomField(models.Model):
+        some_field = MyCustomField()
+
+        class Meta:
+            app_label = "tests"
+
+    with pytest.raises(ConfigError):
+        create_schema(ModelWithCustomField)
+
+    register_field("MyCustomField", int)
+    Schema = create_schema(ModelWithCustomField)
+    print(Schema.json_schema())
+    assert Schema.json_schema()["properties"]["some_field"]["type"] == "integer"