diff --git a/src/marshmallow_generic/decorators.py b/src/marshmallow_generic/decorators.py new file mode 100644 index 0000000..1e46fc5 --- /dev/null +++ b/src/marshmallow_generic/decorators.py @@ -0,0 +1,41 @@ +from collections.abc import Callable +from typing import Any, Optional, TypeVar, overload +from typing_extensions import ParamSpec + +from marshmallow.decorators import post_load as _post_load + + +_R = TypeVar("_R") +_P = ParamSpec("_P") + + +@overload +def post_load( + fn: Callable[_P, _R], + pass_many: bool = False, + pass_original: bool = False, +) -> Callable[_P, _R]: + ... + + +@overload +def post_load( + fn: None = None, + pass_many: bool = False, + pass_original: bool = False, +) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]: + ... + + +def post_load( + fn: Optional[Callable[..., Any]] = None, + pass_many: bool = False, + pass_original: bool = False, +) -> Callable[..., Any]: + """ + Typed overload of the original `marshmallow.post_load` decorator function. + + Generic to ensure that the decorated function retains its type. + Runtime behavior is unchanged. + """ + return _post_load(fn, pass_many=pass_many, pass_original=pass_original) diff --git a/tests/test__util.py b/tests/test__util.py index 7d00205..ee09dfe 100644 --- a/tests/test__util.py +++ b/tests/test__util.py @@ -12,7 +12,9 @@ class GenericInsightMixinTestCase(TestCase): mock_super.return_value = MagicMock(__init_subclass__=mock_super_meth) # Should be `None` by default: - self.assertIsNone(_util.GenericInsightMixin._type_arg) # type: ignore[misc] + self.assertIsNone( + _util.GenericInsightMixin._type_arg # type: ignore[misc] + ) # If the mixin type argument was not specified (still generic), # ensure that the attribute remains `None` on the subclass: diff --git a/tests/test_decorators.py b/tests/test_decorators.py new file mode 100644 index 0000000..e6b7fb6 --- /dev/null +++ b/tests/test_decorators.py @@ -0,0 +1,28 @@ +from collections.abc import Callable +from unittest import TestCase +from unittest.mock import MagicMock, patch + +from marshmallow_generic import decorators + + +class DecoratorsTestCase(TestCase): + @patch.object(decorators, "_post_load") + def test_post_load(self, mock_original_post_load: MagicMock) -> None: + mock_original_post_load.return_value = expected_output = object() + + def test_function(x: int) -> str: + return str(x) + + pass_many, pass_original = MagicMock(), MagicMock() + # Explicit annotation to possibly catch mypy errors: + output: Callable[[int], str] = decorators.post_load( + test_function, + pass_many=pass_many, + pass_original=pass_original, + ) + self.assertIs(expected_output, output) + mock_original_post_load.assert_called_once_with( + test_function, + pass_many=pass_many, + pass_original=pass_original, + )