refactor: add WithState class to be the return value of the store.with_state so that it can report correct signature of its __call__ method

Author: sassanh

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## Upcoming
+
+- refactor: add `WithState` class to be the return value of the `store.with_state` so that it can report correct signature of its `__call__` method
+
 ## Version 0.20.2
 
 - chore(lint): update `ruff` to `0.10.0` and fix linting issues, make `store.subscribe` private

redux/main.py

@@ -59,6 +59,7 @@
 )
 from redux.serialization_mixin import SerializationMixin
 from redux.side_effect_runner import SideEffectRunnerThread
+from redux.with_state import WithState
 
 if TYPE_CHECKING:
     from collections.abc import Awaitable, Callable
@@ -439,15 +440,7 @@ def with_state_decorator(
                 ReturnType,
             ],
         ) -> Callable[Args, ReturnType]:
-            def wrapper(*args: Args.args, **kwargs: Args.kwargs) -> ReturnType:
-                if self._state is None:
-                    msg = 'Store has not been initialized yet.'
-                    raise RuntimeError(msg)
-                return func(selector(self._state), *args, **kwargs)
-
-            wrapper.__name__ = f'with_state:{func.__name__}'
-
-            return wrapper
+            return WithState(store=self, selector=selector, func=cast('Callable', func))
 
         return with_state_decorator
 

redux/with_state.py

@@ -0,0 +1,85 @@
+"""A wrapper for functions that require the current state of the store."""
+
+from __future__ import annotations
+
+import inspect
+from typing import TYPE_CHECKING, Concatenate, Generic
+
+from redux.basic_types import Action, Args, Event, ReturnType, SelectorOutput, State
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from redux.main import Store
+
+
+class WithState(Generic[State, Action, Event, SelectorOutput, ReturnType, Args]):
+    """A wrapper for functions that require the current state of the store."""
+
+    def __init__(
+        self: WithState,
+        *,
+        store: Store[State, Action, Event],
+        selector: Callable[[State], SelectorOutput],
+        func: Callable[
+            Concatenate[SelectorOutput, Args],
+            ReturnType,
+        ],
+    ) -> None:
+        """Initialize the WithState instance."""
+        self._store = store
+        self._selector = selector
+        self._func = func
+        signature = inspect.signature(func)
+        parameters = list(signature.parameters.values())[1:]  # Skip the first parameter
+        if parameters and parameters[0].kind in [
+            inspect.Parameter.POSITIONAL_ONLY,
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+        ]:
+            parameters = parameters[1:]
+        self._signature = signature.replace(parameters=parameters)
+
+    def __call__(
+        self: WithState[
+            State,
+            Action,
+            Event,
+            SelectorOutput,
+            ReturnType,
+            Args,
+        ],
+        *args: Args.args,
+        **kwargs: Args.kwargs,
+    ) -> ReturnType:
+        """Call the wrapped function with the current state of the store."""
+        if self._store._state is None:  # noqa: SLF001
+            msg = 'Store has not been initialized yet.'
+            raise RuntimeError(msg)
+        return self._func(self._selector(self._store._state), *args, **kwargs)  # noqa: SLF001
+
+    def __repr__(
+        self: WithState[
+            State,
+            Action,
+            Event,
+            SelectorOutput,
+            ReturnType,
+            Args,
+        ],
+    ) -> str:
+        """Return the string representation of the WithState instance."""
+        return super().__repr__() + f'(func: {self._func})'
+
+    @property
+    def __signature__(
+        self: WithState[
+            State,
+            Action,
+            Event,
+            SelectorOutput,
+            ReturnType,
+            Args,
+        ],
+    ) -> inspect.Signature:
+        """Get the signature of the wrapped function."""
+        return self._signature

tests/test_with_state.py

@@ -82,8 +82,15 @@ def test_with_state_for_uninitialized_store(
     mocker: MockerFixture,
 ) -> None:
     """Test `with_state` decorator for uninitialized store."""
-    _check = mocker.stub()
-    check = store.with_state(lambda state: state.value)(_check)
+
+    class X:
+        def check(self: X, value: int) -> None:
+            assert value == 0
+
+    instance = X()
+
+    check_spy = mocker.spy(instance, 'check')
+    check = store.with_state(lambda state: state.value)(instance.check)
 
     with pytest.raises(RuntimeError, match=r'^Store has not been initialized yet.$'):
         check()
@@ -94,4 +101,4 @@ def test_with_state_for_uninitialized_store(
 
     store.dispatch(FinishAction())
 
-    _check.assert_called_once_with(0)
+    check_spy.assert_called_once_with(0)