feat: initial implementation of instance level ref cache

jambo/parser/ref_type_parser.py

@@ -1,5 +1,6 @@
 from jambo.exceptions import InternalAssertionException, InvalidSchemaException
 from jambo.parser import GenericTypeParser
+from jambo.types import RefCacheDict
 from jambo.types.json_schema_type import JSONSchema
 from jambo.types.type_parser_options import TypeParserOptions
 
@@ -72,7 +73,7 @@ class RefTypeParser(GenericTypeParser):
         return mapped_type
 
     def _get_ref_from_cache(
-        self, ref_name: str, ref_cache: dict[str, ForwardRef | type | None]
+        self, ref_name: str, ref_cache: RefCacheDict
     ) -> RefType | type | None:
         try:
             ref_state = ref_cache[ref_name]

jambo/schema_converter.py

@@ -1,10 +1,11 @@
 from jambo.exceptions import InvalidSchemaException, UnsupportedSchemaException
 from jambo.parser import ObjectTypeParser, RefTypeParser
-from jambo.types import JSONSchema
+from jambo.types import JSONSchema, RefCacheDict
 
 from jsonschema.exceptions import SchemaError
 from jsonschema.validators import validator_for
 from pydantic import BaseModel
+from typing_extensions import Optional
 
 
 class SchemaConverter:
@@ -16,13 +17,44 @@ class SchemaConverter:
     fields and types. The generated model can be used for data validation and serialization.
     """
 
-    @staticmethod
-    def build(schema: JSONSchema) -> type[BaseModel]:
+    def __init__(self, ref_cache: Optional[RefCacheDict] = None) -> None:
+        if ref_cache is None:
+            ref_cache = dict()
+        self._ref_cache = ref_cache
+
+    def build_with_instance(
+        self,
+        schema: JSONSchema,
+        ref_cache: Optional[RefCacheDict] = None,
+        with_clean_cache: bool = True,
+    ) -> type[BaseModel]:
         """
         Converts a JSON Schema to a Pydantic model.
             :param schema: The JSON Schema to convert.
-        :return: A Pydantic model class.
+            :param ref_cache: An optional reference cache to use during conversion, if provided `with_clean_cache` will be ignored.
+            :param with_clean_cache: Whether to use a clean reference cache for this conversion. Set to True due to API compatibility. Will be set to False in future versions.
+            :return: The generated Pydantic model.
         """
+        if ref_cache is None:
+            ref_cache = self._ref_cache
+
+        return self.build(schema, ref_cache, with_clean_cache)
+
+    @staticmethod
+    def build(
+        schema: JSONSchema,
+        ref_cache: Optional[RefCacheDict] = None,
+        with_clean_cache: bool = True,
+    ) -> type[BaseModel]:
+        """
+        Converts a JSON Schema to a Pydantic model.
+            :param schema: The JSON Schema to convert.
+            :param ref_cache: An optional reference cache to use during conversion, if provided `with_clean_cache` will be ignored.
+            :param with_clean_cache: Whether to use a clean reference cache for this conversion. Set to rue due to API compatibility. Will be set to False in future versions.
+            :return: The generated Pydantic model.
+        """
+        if ref_cache is None or with_clean_cache:
+            ref_cache = dict()
 
         try:
             validator = validator_for(schema)
@@ -46,7 +78,7 @@ class SchemaConverter:
                     schema.get("properties", {}),
                     schema.get("required", []),
                     context=schema,
-                    ref_cache=dict(),
+                    ref_cache=ref_cache,
                     required=True,
                 )
 
@@ -55,7 +87,7 @@ class SchemaConverter:
                     schema["title"],
                     schema,
                     context=schema,
-                    ref_cache=dict(),
+                    ref_cache=ref_cache,
                     required=True,
                 )
                 return parsed_model
@@ -68,6 +100,25 @@ class SchemaConverter:
                     unsupported_field=unsupported_type,
                 )
 
+    def clear_ref_cache(self) -> None:
+        """
+        Clears the reference cache.
+        """
+        self._ref_cache.clear()
+
+    def get_cached_ref(self, ref_name: str):
+        """
+        Gets a cached reference from the reference cache.
+        :param ref_name: The name of the reference to get.
+        :return: The cached reference, or None if not found.
+        """
+        cached_type = self._ref_cache.get(ref_name)
+
+        if isinstance(cached_type, type):
+            return cached_type
+
+        return None
+
     @staticmethod
     def _get_schema_type(schema: JSONSchema) -> str | None:
         """

jambo/types/__init__.py

@@ -4,7 +4,7 @@ from .json_schema_type import (
     JSONSchemaType,
     JSONType,
 )
-from .type_parser_options import TypeParserOptions
+from .type_parser_options import RefCacheDict, TypeParserOptions
 
 
 __all__ = [
@@ -12,5 +12,6 @@ __all__ = [
     "JSONSchemaNativeTypes",
     "JSONType",
     "JSONSchema",
+    "RefCacheDict",
     "TypeParserOptions",
 ]

jambo/types/type_parser_options.py

@@ -1,9 +1,12 @@
 from jambo.types.json_schema_type import JSONSchema
 
-from typing_extensions import ForwardRef, TypedDict
+from typing_extensions import ForwardRef, MutableMapping, TypedDict
+
+
+RefCacheDict = MutableMapping[str, ForwardRef | type | None]
 
 
 class TypeParserOptions(TypedDict):
     required: bool
     context: JSONSchema
-    ref_cache: dict[str, ForwardRef | type | None]
+    ref_cache: RefCacheDict