Postgres ArrayField Inner Type And CharField Choices

ninja/openapi/schema.py

@@ -36,23 +36,21 @@ def __init__(self, api: "NinjaAPI", path_prefix: str) -> None:
         self.securitySchemes: DictStrAny = {}
         self.all_operation_ids: Set = set()
         extra_info = api.openapi_extra.get("info", {})
-        super().__init__(
-            [
-                ("openapi", "3.1.0"),
-                (
-                    "info",
-                    {
-                        "title": api.title,
-                        "version": api.version,
-                        "description": api.description,
-                        **extra_info,
-                    },
-                ),
-                ("paths", self.get_paths()),
-                ("components", self.get_components()),
-                ("servers", api.servers),
-            ]
-        )
+        super().__init__([
+            ("openapi", "3.1.0"),
+            (
+                "info",
+                {
+                    "title": api.title,
+                    "version": api.version,
+                    "description": api.description,
+                    **extra_info,
+                },
+            ),
+            ("paths", self.get_paths()),
+            ("components", self.get_components()),
+            ("servers", api.servers),
+        ])
         for k, v in api.openapi_extra.items():
             if k not in self:
                 self[k] = v
@@ -237,12 +235,10 @@ def _create_multipart_schema_from_models(
         content_type = BODY_CONTENT_TYPES["file"]
 
         # get the various schemas
-        result = merge_schemas(
-            [
-                self._create_schema_from_model(model, remove_level=False)[0]
-                for model in models
-            ]
-        )
+        result = merge_schemas([
+            self._create_schema_from_model(model, remove_level=False)[0]
+            for model in models
+        ])
         result["title"] = "MultiPartBodyParams"
 
         return result, content_type

ninja/orm/fields.py

@@ -1,6 +1,17 @@
 import datetime
 from decimal import Decimal
-from typing import Any, Callable, Dict, List, Tuple, Type, TypeVar, Union, no_type_check
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Literal,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    no_type_check,
+)
 from uuid import UUID
 
 from django.db.models import ManyToManyField
@@ -150,6 +161,15 @@ def get_schema_field(
         internal_type = field.get_internal_type()
         python_type = TYPES[internal_type]
 
+        # Handle CharField with choices
+        if internal_type == "CharField" and field.choices:
+            python_type = Literal[tuple(str(c[0]) for c in field.choices)]
+
+        # Handle PG ArrayField
+        if internal_type == "ArrayField":
+            inner_internal_type, _ = get_schema_field(field.base_field)
+            python_type = python_type[inner_internal_type]
+
         if field.primary_key or blank or null or optional:
             default = None
             nullable = True

ninja/testing/client.py

@@ -127,12 +127,10 @@ def _build_request(
         request.META = request_params.pop("META", {"REMOTE_ADDR": "127.0.0.1"})
         request.FILES = request_params.pop("FILES", {})
 
-        request.META.update(
-            {
-                f"HTTP_{k.replace('-', '_')}": v
-                for k, v in request_params.pop("headers", {}).items()
-            }
-        )
+        request.META.update({
+            f"HTTP_{k.replace('-', '_')}": v
+            for k, v in request_params.pop("headers", {}).items()
+        })
 
         request.headers = HttpHeaders(request.META)
 

tests/test_orm_schemas.py

@@ -1,11 +1,12 @@
-from typing import List
+from typing import List, Literal
 from unittest.mock import Mock
 
 import pytest
 from django.contrib.postgres import fields as ps_fields
 from django.db import models
 from django.db.models import Manager
 
+from ninja import Schema
 from ninja.errors import ConfigError
 from ninja.orm import create_schema
 from ninja.orm.shortcuts import L, S
@@ -43,6 +44,9 @@ class Meta:
 
 def test_all_fields():
     # test all except relational field
+    class TypeChoice(models.TextChoices):
+        NAME_A = "VALUE_A"
+        NAME_B = "VALUE_B"
 
     class AllFields(models.Model):
         bigintegerfield = models.BigIntegerField()
@@ -71,7 +75,11 @@ class AllFields(models.Model):
         timefield = models.TimeField()
         urlfield = models.URLField()
         uuidfield = models.UUIDField()
-        arrayfield = ps_fields.ArrayField(models.CharField())
+        arrayofstringfield = ps_fields.ArrayField(models.CharField())
+        arrayofintegerfield = ps_fields.ArrayField(models.IntegerField())
+        arrayofchoicefield = ps_fields.ArrayField(
+            models.CharField(choices=TypeChoice.choices)
+        )
         cicharfield = ps_fields.CICharField()
         ciemailfield = ps_fields.CIEmailField()
         citextfield = ps_fields.CITextField()
@@ -145,7 +153,24 @@ class Meta:
             "timefield": {"type": "string", "format": "time", "title": "Timefield"},
             "urlfield": {"type": "string", "title": "Urlfield"},
             "uuidfield": {"type": "string", "format": "uuid", "title": "Uuidfield"},
-            "arrayfield": {"type": "array", "items": {}, "title": "Arrayfield"},
+            "arrayofstringfield": {
+                "type": "array",
+                "items": {"type": "string"},
+                "title": "Arrayofstringfield",
+            },
+            "arrayofintegerfield": {
+                "type": "array",
+                "items": {"type": "integer"},
+                "title": "Arrayofintegerfield",
+            },
+            "arrayofchoicefield": {
+                "type": "array",
+                "items": {
+                    "enum": ["VALUE_A", "VALUE_B"],
+                    "type": "string",
+                },
+                "title": "Arrayofchoicefield",
+            },
             "cicharfield": {"type": "string", "title": "Cicharfield"},
             "ciemailfield": {
                 "type": "string",
@@ -182,7 +207,9 @@ class Meta:
             "timefield",
             "urlfield",
             "uuidfield",
-            "arrayfield",
+            "arrayofstringfield",
+            "arrayofintegerfield",
+            "arrayofchoicefield",
             "cicharfield",
             "ciemailfield",
             "citextfield",
@@ -578,3 +605,44 @@ class Meta:
         SomeReqFieldModel, optional_fields=["some_field", "other_field", "optional"]
     )
     assert Schema.json_schema().get("required") is None
+
+
+def test_array_fields():
+    # Test Postgres ArrayField
+
+    class TypeChoice(models.TextChoices):
+        NAME_A = "VALUE_A"
+        NAME_B = "VALUE_B"
+
+    class ArrayModel(models.Model):
+        array_of_int = ps_fields.ArrayField(models.IntegerField())
+        array_of_str = ps_fields.ArrayField(models.CharField(max_length=100))
+        array_of_choice = ps_fields.ArrayField(
+            models.CharField(choices=TypeChoice.choices)
+        )
+
+        class Meta:
+            app_label = "tests"
+
+    class ExpectedSchema(Schema):
+        array_of_int: List[int]
+        array_of_str: List[str]
+        array_of_choice: List[Literal["VALUE_A", "VALUE_B"]]
+
+    model_json_schema = create_schema(ArrayModel).json_schema()
+    expected_json_schema = ExpectedSchema.json_schema()
+
+    assert (
+        model_json_schema["properties"]["array_of_int"]
+        == expected_json_schema["properties"]["array_of_int"]
+    )
+
+    assert (
+        model_json_schema["properties"]["array_of_str"]
+        == expected_json_schema["properties"]["array_of_str"]
+    )
+
+    assert (
+        model_json_schema["properties"]["array_of_choice"]
+        == expected_json_schema["properties"]["array_of_choice"]
+    )