From 760f30d08f31a95a0d6f2a40c335c34c8cc5ae20 Mon Sep 17 00:00:00 2001
From: Vitor Hideyoshi <vitor.h.n.batista@gmail.com>
Date: Thu, 12 Jun 2025 01:54:52 -0300
Subject: [PATCH] Initial Implementation of $ref

---
 jambo/parser/__init__.py             |  2 +
 jambo/parser/ref_type_parser.py      | 72 ++++++++++++++++++++++++++++
 jambo/types/type_parser_options.py   |  4 +-
 tests/parser/test_ref_type_parser.py | 34 +++++++++++++
 4 files changed, 111 insertions(+), 1 deletion(-)
 create mode 100644 jambo/parser/ref_type_parser.py
 create mode 100644 tests/parser/test_ref_type_parser.py

diff --git a/jambo/parser/__init__.py b/jambo/parser/__init__.py
index 86d8f56..b804339 100644
--- a/jambo/parser/__init__.py
+++ b/jambo/parser/__init__.py
@@ -9,6 +9,7 @@ from .boolean_type_parser import BooleanTypeParser
 from .float_type_parser import FloatTypeParser
 from .int_type_parser import IntTypeParser
 from .object_type_parser import ObjectTypeParser
+from .ref_type_parser import RefTypeParser
 from .string_type_parser import StringTypeParser
 
 
@@ -22,4 +23,5 @@ __all__ = [
     "IntTypeParser",
     "ObjectTypeParser",
     "StringTypeParser",
+    "RefTypeParser",
 ]
diff --git a/jambo/parser/ref_type_parser.py b/jambo/parser/ref_type_parser.py
new file mode 100644
index 0000000..c27f62c
--- /dev/null
+++ b/jambo/parser/ref_type_parser.py
@@ -0,0 +1,72 @@
+from jambo.parser import GenericTypeParser
+from jambo.types.type_parser_options import TypeParserOptions
+
+from typing_extensions import Any, ForwardRef, TypeVar, Union, Unpack
+
+
+RefType = TypeVar("RefType", bound=Union[int, str])
+
+
+class RefTypeParser(GenericTypeParser):
+    json_schema_type = "$ref"
+
+    def from_properties_impl(
+        self, name: str, properties: dict[str, Any], **kwargs: Unpack[TypeParserOptions]
+    ) -> tuple[RefType, dict]:
+        if "$ref" not in properties:
+            raise ValueError(f"RefTypeParser: Missing $ref in properties for {name}")
+
+        context = kwargs["context"]
+        required = kwargs.get("required", False)
+
+        if context is None:
+            raise RuntimeError(
+                f"RefTypeParser: Missing $content in properties for {name}"
+            )
+
+        if not properties["$ref"].startswith("#"):
+            raise ValueError(
+                "At the moment, only local references are supported. "
+                "Look into $defs and # for recursive references."
+            )
+
+        ref_type = None
+        mapped_properties = {}
+
+        if properties["$ref"] == "#":
+            if "title" not in context:
+                raise ValueError(
+                    "RefTypeParser: Missing title in properties for $ref #"
+                )
+
+            ref_type = ForwardRef(context["title"])
+
+        elif properties["$ref"].startswith("#/$defs/"):
+            target_name = None
+            target_property = context
+            for prop_name in properties["$ref"].split("/")[1:]:
+                if prop_name not in target_property:
+                    raise ValueError(
+                        f"RefTypeParser: Missing {prop_name} in"
+                        " properties for $ref {properties['$ref']}"
+                    )
+                target_name = prop_name
+                target_property = target_property[prop_name]
+
+            if target_name is None or target_property is None:
+                raise ValueError(f"RefTypeParser: Invalid $ref {properties['$ref']}")
+
+            ref_type, mapped_properties = GenericTypeParser.type_from_properties(
+                target_name, target_property, **kwargs
+            )
+
+        else:
+            raise ValueError(
+                "RefTypeParser: Invalid $ref format. "
+                "Only local references are supported."
+            )
+
+        if not required:
+            mapped_properties["default"] = None
+
+        return ref_type, mapped_properties
diff --git a/jambo/types/type_parser_options.py b/jambo/types/type_parser_options.py
index b75490e..ae96338 100644
--- a/jambo/types/type_parser_options.py
+++ b/jambo/types/type_parser_options.py
@@ -1,5 +1,7 @@
-from typing_extensions import TypedDict
+from typing_extensions import Any, NotRequired, TypedDict
 
 
 class TypeParserOptions(TypedDict):
     required: bool
+    context: dict[str, Any]
+    ref_cache: NotRequired[dict[str, type]]
diff --git a/tests/parser/test_ref_type_parser.py b/tests/parser/test_ref_type_parser.py
new file mode 100644
index 0000000..5ab96f5
--- /dev/null
+++ b/tests/parser/test_ref_type_parser.py
@@ -0,0 +1,34 @@
+from jambo.parser import RefTypeParser
+
+from unittest import TestCase
+
+
+class TestRefTypeParser(TestCase):
+    def test_ref_type_parser_local_ref(self):
+        properties = {
+            "title": "person",
+            "$ref": "#/$defs/person",
+            "$defs": {
+                "person": {
+                    "type": "object",
+                    "properties": {
+                        "name": {"type": "string"},
+                        "age": {"type": "integer"},
+                    },
+                }
+            },
+        }
+
+        type_parsing, type_validator = RefTypeParser().from_properties(
+            properties=properties,
+            name="placeholder",
+            context=properties,
+            required=True,
+        )
+
+        self.assertIsInstance(type_parsing, type)
+
+        obj = type_parsing(name="John", age=30)
+
+        self.assertEqual(obj.name, "John")
+        self.assertEqual(obj.age, 30)