Make JSONField support type annotation and OpanAPI document generation (

#1763) * Make JSONField support type annotation and OpanAPI document generation fix: codacy * fix: pass test and format code * test: add test case for json field pydantic type * fix: input * test: add testcase for init with pydantic type * chore(docs): add changelog * fix: codacy issue * chore(docs): add docs to code * refactor: use model_dump instead of dict method * fix: pass lint * chore: make style * fix: test python version is before 3.10 * fix: pytest warning * fix: codacy
tortoise · Nov 18, 2024 · defd780 · defd780
1 parent ea34c3d
commit defd780
Show file tree

Hide file tree

Showing 11 changed files with 356 additions and 309 deletions.
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -8,6 +8,18 @@ Changelog
 
 0.21
 ====
+0.21.8
+------
+Fixed
+^^^^^
+- TODO
+
+Added
+^^^^^
+- JSONField adds optional generic support, and supports OpenAPI document generation by specifying `field_type` as a pydantic BaseModel (#1763)
+
+
+
 0.21.7
 ------
 Fixed

diff --git a/CONTRIBUTORS.rst b/CONTRIBUTORS.rst
@@ -61,6 +61,7 @@ Contributors
 * Andrea Magistà ``@vlakius``
 * Daniel Szucs ``@Quasar6X``
 * Rui Catarino ``@ruitcatarino``
+* Lance Moe ``@lancemoe``
 
 Special Thanks
 ==============

diff --git a/docs/query.rst b/docs/query.rst
@@ -243,7 +243,7 @@ In PostgreSQL and MYSQL, you can use the ``contains``, ``contained_by`` and ``fi
 .. code-block:: python3
 
     class JSONModel:
-        data = fields.JSONField()
+        data = fields.JSONField[list]()
 
     await JSONModel.create(data=["text", 3, {"msg": "msg2"}])
     obj = await JSONModel.filter(data__contains=[{"msg": "msg2"}]).first()
@@ -257,7 +257,7 @@ In PostgreSQL and MYSQL, you can use the ``contains``, ``contained_by`` and ``fi
 .. code-block:: python3
 
     class JSONModel:
-        data = fields.JSONField()
+        data = fields.JSONField[dict]()
 
     await JSONModel.create(data={"breed": "labrador",
                                  "owner": {

diff --git a/examples/postgres.py b/examples/postgres.py
@@ -8,7 +8,7 @@
 
 class Report(Model):
     id = fields.IntField(primary_key=True)
-    content = fields.JSONField()
+    content = fields.JSONField[dict]()
 
     def __str__(self):
         return str(self.id)

diff --git a/tests/contrib/test_pydantic.py b/tests/contrib/test_pydantic.py
diff --git a/tests/test_queryset_reuse.py b/tests/test_queryset_reuse.py
@@ -1,7 +1,4 @@
-from tests.testmodels import (
-    Event,
-    Tournament,
-)
+from tests.testmodels import Event, Tournament
 from tortoise.contrib import test
 from tortoise.contrib.test.condition import NotEQ
 from tortoise.expressions import F

diff --git a/tests/testmodels.py b/tests/testmodels.py
@@ -12,7 +12,7 @@
 from typing import List, Union
 
 import pytz
-from pydantic import ConfigDict
+from pydantic import BaseModel, ConfigDict
 
 from tortoise import fields
 from tortoise.exceptions import ValidationError
@@ -34,6 +34,15 @@ def generate_token():
     return binascii.hexlify(os.urandom(16)).decode("ascii")
 
 
+class TestSchemaForJSONField(BaseModel):
+    foo: int
+    bar: str
+    __test__ = False
+
+
+json_pydantic_default = TestSchemaForJSONField(foo=1, bar="baz")
+
+
 class Author(Model):
     name = fields.CharField(max_length=255)
 
@@ -286,21 +295,30 @@ class FloatFields(Model):
     floatnum_null = fields.FloatField(null=True)
 
 
+def raise_if_not_dict_or_list(value: Union[dict, list]):
+    if not isinstance(value, (dict, list)):
+        raise ValidationError("Value must be a dict or list.")
+
+
 class JSONFields(Model):
     """
     This model contains many JSON blobs
     """
 
-    @staticmethod
-    def dict_or_list(value: Union[dict, list]):
-        if not isinstance(value, (dict, list)):
-            raise ValidationError("Value must be a dict or list.")
-
     id = fields.IntField(primary_key=True)
-    data = fields.JSONField()
-    data_null = fields.JSONField(null=True)
-    data_default = fields.JSONField(default={"a": 1})
-    data_validate = fields.JSONField(null=True, validators=[lambda v: JSONFields.dict_or_list(v)])
+    data = fields.JSONField()  # type: ignore # Test cases where generics are not provided
+    data_null = fields.JSONField[Union[dict, list]](null=True)
+    data_default = fields.JSONField[dict](default={"a": 1})
+
+    # From Python 3.10 onwards, validator can be defined with staticmethod
+    data_validate = fields.JSONField[Union[dict, list]](
+        null=True, validators=[raise_if_not_dict_or_list]
+    )
+
+    # Test cases where generics are provided and the type is a pydantic base model
+    data_pydantic = fields.JSONField[TestSchemaForJSONField](
+        default=json_pydantic_default, field_type=TestSchemaForJSONField
+    )
 
 
 class UUIDFields(Model):

diff --git a/tests/utils/test_describe_model.py b/tests/utils/test_describe_model.py
@@ -9,11 +9,13 @@
     SourceFields,
     StraightFields,
     Team,
+    TestSchemaForJSONField,
     Tournament,
     UUIDFkRelatedModel,
     UUIDFkRelatedNullModel,
     UUIDM2MRelatedModel,
     UUIDPkModel,
+    json_pydantic_default,
 )
 from tortoise import Tortoise, fields
 from tortoise.contrib import test
@@ -1392,6 +1394,26 @@ def test_describe_model_json(self):
                         "docstring": None,
                         "constraints": {},
                     },
+                    {
+                        "name": "data_pydantic",
+                        "field_type": "JSONField",
+                        "db_column": "data_pydantic",
+                        "db_field_types": {
+                            "": "JSON",
+                            "mssql": "NVARCHAR(MAX)",
+                            "oracle": "NCLOB",
+                            "postgres": "JSONB",
+                        },
+                        "python_type": "tests.testmodels.TestSchemaForJSONField",
+                        "generated": False,
+                        "nullable": False,
+                        "unique": False,
+                        "indexed": False,
+                        "default": "foo=1 bar='baz'",
+                        "description": None,
+                        "docstring": None,
+                        "constraints": {},
+                    },
                 ],
                 "fk_fields": [],
                 "backward_fk_fields": [],
@@ -1511,6 +1533,26 @@ def test_describe_model_json_native(self):
                         "docstring": None,
                         "constraints": {},
                     },
+                    {
+                        "name": "data_pydantic",
+                        "field_type": fields.JSONField,
+                        "db_column": "data_pydantic",
+                        "db_field_types": {
+                            "": "JSON",
+                            "mssql": "NVARCHAR(MAX)",
+                            "oracle": "NCLOB",
+                            "postgres": "JSONB",
+                        },
+                        "python_type": TestSchemaForJSONField,
+                        "generated": False,
+                        "nullable": False,
+                        "unique": False,
+                        "indexed": False,
+                        "default": json_pydantic_default,
+                        "description": None,
+                        "docstring": None,
+                        "constraints": {},
+                    },
                 ],
                 "fk_fields": [],
                 "backward_fk_fields": [],