fix: schema class can set Meta.unknown

Signed-off-by: Daniel Bluhm <dbluhm@pm.me>

Author: dbluhm

aries_cloudagent/messaging/models/base.py

@@ -5,7 +5,7 @@
 
 from abc import ABC
 from collections import namedtuple
-from typing import Mapping, Union
+from typing import Literal, Mapping, Optional, Type, TypeVar, Union, cast, overload
 
 from marshmallow import Schema, post_dump, pre_load, post_load, ValidationError, EXCLUDE
 
@@ -17,7 +17,7 @@
 SerDe = namedtuple("SerDe", "ser de")
 
 
-def resolve_class(the_cls, relative_cls: type = None):
+def resolve_class(the_cls, relative_cls: Optional[type] = None) -> type:
     """
     Resolve a class.
 
@@ -38,6 +38,10 @@ def resolve_class(the_cls, relative_cls: type = None):
     elif isinstance(the_cls, str):
         default_module = relative_cls and relative_cls.__module__
         resolved = ClassLoader.load_class(the_cls, default_module)
+    else:
+        raise TypeError(
+            f"Could not resolve class from {the_cls}; incorrect type {type(the_cls)}"
+        )
     return resolved
 
 
@@ -70,6 +74,9 @@ class BaseModelError(BaseError):
     """Base exception class for base model errors."""
 
 
+ModelType = TypeVar("ModelType", bound="BaseModel")
+
+
 class BaseModel(ABC):
     """Base model that provides convenience methods."""
 
@@ -94,18 +101,24 @@ def __init__(self):
             )
 
     @classmethod
-    def _get_schema_class(cls):
+    def _get_schema_class(cls) -> Type["BaseModelSchema"]:
         """
         Get the schema class.
 
         Returns:
             The resolved schema class
 
         """
-        return resolve_class(cls.Meta.schema_class, cls)
+        resolved = resolve_class(cls.Meta.schema_class, cls)
+        if issubclass(resolved, BaseModelSchema):
+            return resolved
+
+        raise TypeError(
+            f"Resolved class is not a subclass of BaseModelSchema: {resolved}"
+        )
 
     @property
-    def Schema(self) -> type:
+    def Schema(self) -> Type["BaseModelSchema"]:
         """
         Accessor for the model's schema class.
 
@@ -115,8 +128,46 @@ def Schema(self) -> type:
         """
         return self._get_schema_class()
 
+    @overload
     @classmethod
-    def deserialize(cls, obj, unknown: str = None, none2none: str = False):
+    def deserialize(
+        cls: Type[ModelType],
+        obj,
+        *,
+        unknown: Optional[str] = None,
+    ) -> ModelType:
+        ...
+
+    @overload
+    @classmethod
+    def deserialize(
+        cls: Type[ModelType],
+        obj,
+        *,
+        none2none: Literal[False],
+        unknown: Optional[str] = None,
+    ) -> ModelType:
+        ...
+
+    @overload
+    @classmethod
+    def deserialize(
+        cls: Type[ModelType],
+        obj,
+        *,
+        none2none: Literal[True],
+        unknown: Optional[str] = None,
+    ) -> Optional[ModelType]:
+        ...
+
+    @classmethod
+    def deserialize(
+        cls: Type[ModelType],
+        obj,
+        *,
+        unknown: Optional[str] = None,
+        none2none: bool = False,
+    ) -> Optional[ModelType]:
         """
         Convert from JSON representation to a model instance.
 
@@ -132,18 +183,41 @@ def deserialize(cls, obj, unknown: str = None, none2none: str = False):
         if obj is None and none2none:
             return None
 
-        schema = cls._get_schema_class()(unknown=unknown or EXCLUDE)
+        schema_cls = cls._get_schema_class()
+        schema = schema_cls(unknown=unknown or schema_cls.Meta.unknown)
+
         try:
-            return schema.loads(obj) if isinstance(obj, str) else schema.load(obj)
+            return cast(
+                ModelType,
+                schema.loads(obj) if isinstance(obj, str) else schema.load(obj),
+            )
         except (AttributeError, ValidationError) as err:
             LOGGER.exception(f"{cls.__name__} message validation error:")
             raise BaseModelError(f"{cls.__name__} schema validation failed") from err
 
+    @overload
     def serialize(
         self,
-        as_string=False,
-        unknown: str = None,
+        *,
+        as_string: Literal[True],
+        unknown: Optional[str] = None,
+    ) -> str:
+        ...
+
+    @overload
+    def serialize(
+        self,
+        *,
+        unknown: Optional[str] = None,
     ) -> dict:
+        ...
+
+    def serialize(
+        self,
+        *,
+        as_string: bool = False,
+        unknown: Optional[str] = None,
+    ) -> Union[str, dict]:
         """
         Create a JSON-compatible dict representation of the model instance.
 
@@ -154,7 +228,8 @@ def serialize(
             A dict representation of this model, or a JSON string if as_string is True
 
         """
-        schema = self.Schema(unknown=unknown or EXCLUDE)
+        schema_cls = self._get_schema_class()
+        schema = schema_cls(unknown=unknown or schema_cls.Meta.unknown)
         try:
             return (
                 schema.dumps(self, separators=(",", ":"))
@@ -168,18 +243,17 @@ def serialize(
             ) from err
 
     @classmethod
-    def serde(cls, obj: Union["BaseModel", Mapping]) -> SerDe:
+    def serde(cls, obj: Union["BaseModel", Mapping]) -> Optional[SerDe]:
         """Return serialized, deserialized representations of input object."""
+        if obj is None:
+            return None
 
-        return (
-            SerDe(obj.serialize(), obj)
-            if isinstance(obj, BaseModel)
-            else None
-            if obj is None
-            else SerDe(obj, cls.deserialize(obj))
-        )
+        if isinstance(obj, BaseModel):
+            return SerDe(obj.serialize(), obj)
+
+        return SerDe(obj, cls.deserialize(obj))
 
-    def validate(self, unknown: str = None):
+    def validate(self, unknown: Optional[str] = None):
         """Validate a constructed model."""
         schema = self.Schema(unknown=unknown)
         errors = schema.validate(self.serialize())
@@ -191,7 +265,7 @@ def validate(self, unknown: str = None):
     def from_json(
         cls,
         json_repr: Union[str, bytes],
-        unknown: str = None,
+        unknown: Optional[str] = None,
     ):
         """
         Parse a JSON string into a model instance.
@@ -218,7 +292,7 @@ def to_json(self, unknown: str = None) -> str:
             A JSON representation of this message
 
         """
-        return json.dumps(self.serialize(unknown=unknown or EXCLUDE))
+        return json.dumps(self.serialize(unknown=unknown))
 
     def __repr__(self) -> str:
         """
@@ -246,6 +320,7 @@ class Meta:
         model_class = None
         skip_values = [None]
         ordered = True
+        unknown = EXCLUDE
 
     def __init__(self, *args, **kwargs):
         """

aries_cloudagent/utils/classloader.py

@@ -7,7 +7,7 @@
 from importlib import import_module
 from importlib.util import find_spec, resolve_name
 from types import ModuleType
-from typing import Sequence, Type
+from typing import Optional, Sequence, Type
 
 from ..core.error import BaseError
 
@@ -75,7 +75,10 @@ def load_module(cls, mod_path: str, package: str = None) -> ModuleType:
 
     @classmethod
     def load_class(
-        cls, class_name: str, default_module: str = None, package: str = None
+        cls,
+        class_name: str,
+        default_module: Optional[str] = None,
+        package: Optional[str] = None,
     ):
         """
         Resolve a complete class path (ie. typing.Dict) to the class itself.