add register_type public function (#1352)

* register_field utility

* mypy lint

Author: vitalik

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")

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"]

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

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"