From 39e1ac3390ddfdec59dae7cfd3a3c724cee51e08 Mon Sep 17 00:00:00 2001
From: Daniil Fajnberg <mail@daniil.fajnberg.de>
Date: Fri, 10 Mar 2023 13:04:00 +0100
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20Overload=20`@post=5Fload`=20decorat?=
 =?UTF-8?q?or=20to=20retain=20the=20function=20type?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/marshmallow_generic/decorators.py | 41 +++++++++++++++++++++++++++
 tests/test__util.py                   |  4 ++-
 tests/test_decorators.py              | 28 ++++++++++++++++++
 3 files changed, 72 insertions(+), 1 deletion(-)
 create mode 100644 src/marshmallow_generic/decorators.py
 create mode 100644 tests/test_decorators.py

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,
+        )