refactor: disuse lru_cache in favor of an _LRUDict

Author: dairiki

marshmallow_dataclass/__init__.py

@@ -42,7 +42,7 @@ class User:
 import types
 import warnings
 from enum import Enum
-from functools import lru_cache, partial
+from functools import partial
 from typing import (
     Any,
     Callable,
@@ -106,8 +106,12 @@ def get_origin(tp):
 
 
 if sys.version_info >= (3, 7):
+    from typing import OrderedDict
+
     TypeVar_ = TypeVar
 else:
+    from typing_extensions import OrderedDict
+
     TypeVar_ = type
 
 if sys.version_info >= (3, 10):
@@ -118,6 +122,7 @@ def get_origin(tp):
 
 NoneType = type(None)
 _U = TypeVar("_U")
+_V = TypeVar("_V")
 
 # Whitelist of dataclass members that will be copied to generated schema.
 MEMBERS_WHITELIST: Set[str] = {"Meta"}
@@ -487,13 +492,59 @@ def class_schema(
             clazz_frame = _maybe_get_callers_frame(clazz)
         if clazz_frame is not None:
             localns = clazz_frame.f_locals
-    with _SchemaContext(globalns, localns):
-        schema = _internal_class_schema(clazz, base_schema)
+
+    if base_schema is None:
+        base_schema = marshmallow.Schema
+
+    with _SchemaContext(globalns, localns, base_schema):
+        schema = _internal_class_schema(clazz)
 
     assert not isinstance(schema, _Future)
     return schema
 
 
+class _LRUDict(OrderedDict[_U, _V]):
+    """Limited-length dict which discards LRU entries."""
+
+    def __init__(self, maxsize: int = 128):
+        self.maxsize = maxsize
+        super().__init__()
+
+    def __setitem__(self, key: _U, value: _V) -> None:
+        super().__setitem__(key, value)
+        super().move_to_end(key)
+
+        while len(self) > self.maxsize:
+            oldkey = next(iter(self))
+            super().__delitem__(oldkey)
+
+    def __getitem__(self, key: _U) -> _V:
+        val = super().__getitem__(key)
+        super().move_to_end(key)
+        return val
+
+    _T = TypeVar("_T")
+
+    @overload
+    def get(self, key: _U) -> Optional[_V]:
+        ...
+
+    @overload
+    def get(self, key: _U, default: _T) -> Union[_V, _T]:
+        ...
+
+    def get(self, key: _U, default: Any = None) -> Any:
+        try:
+            return self.__getitem__(key)
+        except KeyError:
+            return default
+
+
+_schema_cache = _LRUDict[Hashable, Type[marshmallow.Schema]](
+    MAX_CLASS_SCHEMA_CACHE_SIZE
+)
+
+
 class InvalidStateError(Exception):
     """Raised when an operation is performed on a future that is not
     allowed in the current state.
@@ -623,7 +674,7 @@ class _SchemaContext:
 
     globalns: Optional[Dict[str, Any]] = None
     localns: Optional[Dict[str, Any]] = None
-    base_schema: Optional[Type[marshmallow.Schema]] = None
+    base_schema: Type[marshmallow.Schema] = marshmallow.Schema
     generic_args: Optional[_GenericArgs] = None
     seen_classes: Dict[type, _Future[Type[marshmallow.Schema]]] = dataclasses.field(
         default_factory=dict
@@ -633,12 +684,9 @@ def get_type_mapping(
         self, include_marshmallow_default: bool = False
     ) -> _TypeMapping:
         default_mapping = marshmallow.Schema.TYPE_MAPPING
-        if self.base_schema is not None:
-            mappings = [self.base_schema.TYPE_MAPPING]
-            if include_marshmallow_default:
-                mappings.append(default_mapping)
-        else:
-            mappings = [default_mapping]
+        mappings = [self.base_schema.TYPE_MAPPING]
+        if include_marshmallow_default and mappings[0] is not default_mapping:
+            mappings.append(default_mapping)
         return _TypeMapping(*mappings)
 
     def __enter__(self) -> "_SchemaContext":
@@ -672,15 +720,19 @@ def top(self) -> _U:
 _schema_ctx_stack = _LocalStack[_SchemaContext]()
 
 
-@lru_cache(maxsize=MAX_CLASS_SCHEMA_CACHE_SIZE)
 def _internal_class_schema(
     clazz: type,
-    base_schema: Optional[Type[marshmallow.Schema]] = None,
 ) -> Union[Type[marshmallow.Schema], _Future[Type[marshmallow.Schema]]]:
     schema_ctx = _schema_ctx_stack.top
     if clazz in schema_ctx.seen_classes:
         return schema_ctx.seen_classes[clazz]
 
+    cache_key = clazz, schema_ctx.base_schema
+    try:
+        return _schema_cache[cache_key]
+    except KeyError:
+        pass
+
     future: _Future[Type[marshmallow.Schema]] = _Future()
     schema_ctx.seen_classes[clazz] = future
 
@@ -721,9 +773,7 @@ def _internal_class_schema(
     type_hints = get_type_hints(
         clazz, globalns=schema_ctx.globalns, localns=schema_ctx.localns
     )
-    with dataclasses.replace(
-        schema_ctx, base_schema=base_schema, generic_args=generic_args
-    ):
+    with dataclasses.replace(schema_ctx, generic_args=generic_args):
         attributes.update(
             (
                 field.name,
@@ -738,9 +788,10 @@ def _internal_class_schema(
         )
 
     schema_class: Type[marshmallow.Schema] = type(
-        clazz.__name__, (_base_schema(clazz, base_schema),), attributes
+        clazz.__name__, (_base_schema(clazz, schema_ctx.base_schema),), attributes
     )
     future.set_result(schema_class)
+    _schema_cache[cache_key] = schema_class
     return schema_class
 
 
@@ -958,8 +1009,7 @@ def _field_for_dataclass(
         nested = typ.Schema
     else:
         assert isinstance(typ, Hashable)
-        schema_ctx = _schema_ctx_stack.top
-        nested = _internal_class_schema(typ, schema_ctx.base_schema)
+        nested = _internal_class_schema(typ)  # type: ignore[arg-type] # FIXME
         if isinstance(nested, _Future):
             nested = nested.result
 
@@ -994,6 +1044,8 @@ def field_for_schema(
     >>> field_for_schema(str, metadata={"marshmallow_field": marshmallow.fields.Url()}).__class__
     <class 'marshmallow.fields.Url'>
     """
+    if base_schema is None:
+        base_schema = marshmallow.Schema
     localns = typ_frame.f_locals if typ_frame is not None else None
     with _SchemaContext(localns=localns, base_schema=base_schema):
         return _field_for_schema(typ, default, metadata)

tests/test_lrudict.py

@@ -0,0 +1,35 @@
+from marshmallow_dataclass import _LRUDict
+
+
+def test_LRUDict_getitem_moves_to_end() -> None:
+    d = _LRUDict[str, str]()
+    d["a"] = "aval"
+    d["b"] = "bval"
+    assert list(d.items()) == [("a", "aval"), ("b", "bval")]
+    assert d["a"] == "aval"
+    assert list(d.items()) == [("b", "bval"), ("a", "aval")]
+
+
+def test_LRUDict_get_moves_to_end() -> None:
+    d = _LRUDict[str, str]()
+    d["a"] = "aval"
+    d["b"] = "bval"
+    assert list(d.items()) == [("a", "aval"), ("b", "bval")]
+    assert d.get("a") == "aval"
+    assert list(d.items()) == [("b", "bval"), ("a", "aval")]
+
+
+def test_LRUDict_setitem_moves_to_end() -> None:
+    d = _LRUDict[str, str]()
+    d["a"] = "aval"
+    d["b"] = "bval"
+    assert list(d.items()) == [("a", "aval"), ("b", "bval")]
+    d["a"] = "newval"
+    assert list(d.items()) == [("b", "bval"), ("a", "newval")]
+
+
+def test_LRUDict_discards_oldest() -> None:
+    d = _LRUDict[str, str](maxsize=1)
+    d["a"] = "aval"
+    d["b"] = "bval"
+    assert list(d.items()) == [("b", "bval")]