diff --git a/src/marshmallow_generic/schema.py b/src/marshmallow_generic/schema.py index c857bc0..797faac 100644 --- a/src/marshmallow_generic/schema.py +++ b/src/marshmallow_generic/schema.py @@ -66,3 +66,44 @@ class GenericSchema(GenericInsightMixin[_T], Schema): correctly based on the type argument passed to a specific subclass. """ ... + + @overload # type: ignore[override] + def loads( + self, + json_data: str, + *, + many: Literal[True], + partial: Union[bool, Sequence[str], set[str], None] = None, + unknown: Optional[str] = None, + **kwargs: Any, + ) -> list[_T]: + ... + + @overload + def loads( + self, + json_data: str, + *, + many: Optional[Literal[False]] = None, + partial: Union[bool, Sequence[str], set[str], None] = None, + unknown: Optional[str] = None, + **kwargs: Any, + ) -> _T: + ... + + def loads( + self, + json_data: str, + *, + many: Optional[bool] = None, + partial: Union[bool, Sequence[str], set[str], None] = None, + unknown: Optional[str] = None, + **kwargs: Any, + ) -> Union[list[_T], _T]: + """ + Same as `marshmallow.Schema.loads` at runtime. + + Annotations ensure that type checkers will infer the return type + correctly based on the type argument passed to a specific subclass. + """ + ... diff --git a/tests/test_schema.py b/tests/test_schema.py index 0c4f6e2..59e17e6 100644 --- a/tests/test_schema.py +++ b/tests/test_schema.py @@ -20,7 +20,7 @@ class GenericSchemaTestCase(TestCase): mock__get_type_arg.assert_called_once_with() mock_cls.assert_called_once_with(**mock_data) - def test_load(self) -> None: + def test_load_and_loads(self) -> None: """Mainly for static type checking purposes.""" class Foo: @@ -29,9 +29,16 @@ class GenericSchemaTestCase(TestCase): class TestSchema(schema.GenericSchema[Foo]): pass - single: Foo = TestSchema().load({}) + single: Foo + single = TestSchema().load({}) + self.assertIsInstance(single, Foo) + single = TestSchema().loads("{}") self.assertIsInstance(single, Foo) - multiple: list[Foo] = TestSchema().load([{}], many=True) + multiple: list[Foo] + multiple = TestSchema().load([{}], many=True) + self.assertIsInstance(multiple, list) + self.assertIsInstance(multiple[0], Foo) + multiple = TestSchema().loads("[{}]", many=True) self.assertIsInstance(multiple, list) self.assertIsInstance(multiple[0], Foo)