From dee3fa77836f88c2a8566c0b3eef4cdd0ceca708 Mon Sep 17 00:00:00 2001
From: Jose Tomas Robles Hahn <jtrobles@cordada.com>
Date: Fri, 25 Oct 2024 20:33:06 -0300
Subject: [PATCH] fix(extras): Fix generation of JSON Schema for Pydantic `Rut`
 type
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Attempting to generate a JSON Schema using
`pydantic.BaseModel.model_json_schema()` or
`pydantic.TypeAdapter.json_schema()` fails with the following error message:

> pydantic.errors.PydanticInvalidForJsonSchema:
>   Cannot generate a JsonSchema for core_schema.PlainValidatorFunctionSchema ({
>     'type': 'no-info',
>     'function': <function _RutPydanticAnnotation.__get_pydantic_core_schema__.<locals>.validate_from_str at 0x…>
>   })

This commit fixes the above error by implementing `__get_pydantic_json_schema__`
to override the generated JSON Schema.
---
 src/cl_sii/extras/pydantic_types.py     | 19 ++++++++++++++-
 src/tests/test_extras_pydantic_types.py | 32 +++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/src/cl_sii/extras/pydantic_types.py b/src/cl_sii/extras/pydantic_types.py
index cd0f3a89..b1dd73e7 100644
--- a/src/cl_sii/extras/pydantic_types.py
+++ b/src/cl_sii/extras/pydantic_types.py
@@ -41,6 +41,9 @@ class _RutPydanticAnnotation:
         - Customizing the core schema and JSON schema:
             https://docs.pydantic.dev/2.9/architecture/#customizing-the-core-schema-and-json-schema
             (https://github.com/pydantic/pydantic/blob/v2.9.2/docs/architecture.md#customizing-the-core-schema-and-json-schema)
+        - Implementing __get_pydantic_json_schema__:
+            https://docs.pydantic.dev/2.9/concepts/json_schema/#implementing-__get_pydantic_json_schema__
+            (https://github.com/pydantic/pydantic/blob/v2.9.2/docs/concepts/json_schema.md#implementing-__get_pydantic_json_schema__-)
 
     Examples:
 
@@ -73,6 +76,7 @@ class _RutPydanticAnnotation:
     '78773510-K'
     >>> example_type_adapter.dump_json(cl_sii.rut.Rut('78773510-K'))
     b'"78773510-K"'
+    >>> example_json_schema = example_type_adapter.json_schema()
     """
 
     RUT_CANONICAL_STRICT_REGEX: ClassVar[Pattern] = re.compile(
@@ -99,7 +103,7 @@ def validate_from_str(value: str) -> cl_sii.rut.Rut:
 
         from_str_schema = pydantic_core.core_schema.chain_schema(
             [
-                pydantic_core.core_schema.str_schema(pattern=cls.RUT_CANONICAL_STRICT_REGEX),
+                cls.str_schema(),
                 pydantic_core.core_schema.no_info_plain_validator_function(validate_from_str),
             ]
         )
@@ -117,6 +121,19 @@ def validate_from_str(value: str) -> cl_sii.rut.Rut:
             ),
         )
 
+    @classmethod
+    def __get_pydantic_json_schema__(
+        cls,
+        core_schema: pydantic_core.core_schema.CoreSchema,
+        handler: pydantic.GetJsonSchemaHandler,
+    ) -> pydantic.json_schema.JsonSchemaValue:
+        core_schema = cls.str_schema()
+        return handler(core_schema)
+
+    @classmethod
+    def str_schema(cls) -> pydantic_core.core_schema.CoreSchema:
+        return pydantic_core.core_schema.str_schema(pattern=cls.RUT_CANONICAL_STRICT_REGEX)
+
 
 Rut = Annotated[cl_sii.rut.Rut, _RutPydanticAnnotation]
 """
diff --git a/src/tests/test_extras_pydantic_types.py b/src/tests/test_extras_pydantic_types.py
index e66221d8..3664e122 100644
--- a/src/tests/test_extras_pydantic_types.py
+++ b/src/tests/test_extras_pydantic_types.py
@@ -133,3 +133,35 @@ def test_deserialize_invalid(self) -> None:
 
                 with self.assertRaises(pydantic.ValidationError):
                     self.pydantic_type_adapter.validate_json(data)
+
+    def test_json_schema_for_validation(self) -> None:
+        # -----Arrange-----
+
+        expected_json_schema = {
+            'type': 'string',
+            'pattern': '^(\\d{1,8})-([\\dK])$',
+        }
+
+        # -----Act-----
+
+        actual_json_schema = self.pydantic_type_adapter.json_schema(mode='validation')
+
+        # -----Assert-----
+
+        self.assertEqual(expected_json_schema, actual_json_schema)
+
+    def test_json_schema_for_serialization(self) -> None:
+        # -----Arrange-----
+
+        expected_json_schema = {
+            'type': 'string',
+            'pattern': '^(\\d{1,8})-([\\dK])$',
+        }
+
+        # -----Act-----
+
+        actual_json_schema = self.pydantic_type_adapter.json_schema(mode='serialization')
+
+        # -----Assert-----
+
+        self.assertEqual(expected_json_schema, actual_json_schema)