feat: support with_state to be applied to methods of classes, not just functions

Author: sassanh

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 - chore: add badges in `README.md` and classifiers in `pyproject.toml`
 - refactor: move the common code for manipulating the signature of the wrapped functions in `WithStore` and `Autorun` to a utility function
+- feat: support `with_state` to be applied to methods of classes, not just functions
 
 ## Version 0.23.0
 

redux/basic_types.py

@@ -361,15 +361,25 @@ def __call__(
 # With Store
 
 
+Self = TypeVar('Self', bound=object, infer_variance=True)
+
+
 class WithStateDecorator(
     Protocol,
     Generic[SelectorOutput],
 ):
+    @overload
     def __call__(
         self: WithStateDecorator,
         func: Callable[Concatenate[SelectorOutput, Args], ReturnType],
     ) -> Callable[Args, ReturnType]: ...
 
+    @overload
+    def __call__(
+        self: WithStateDecorator,
+        func: Callable[Concatenate[Self, SelectorOutput, Args], ReturnType],
+    ) -> Callable[Concatenate[Self, Args], ReturnType]: ...
+
 
 class EventSubscriber(Protocol):
     def __call__(

redux/main.py

@@ -42,6 +42,7 @@
     ReducerType,
     ReturnType,
     SelectorOutput,
+    Self,
     SnapshotAtom,
     State,
     StoreOptions,
@@ -384,20 +385,67 @@ def with_state(
         `store._state` is also possible.
         """
 
+        @overload
+        def with_state_decorator(
+            func: Callable[
+                Concatenate[SelectorOutput, Args],
+                ReturnType,
+            ],
+        ) -> Callable[Args, ReturnType]: ...
+        @overload
+        def with_state_decorator(
+            func: Callable[
+                Concatenate[Self, SelectorOutput, Args],
+                ReturnType,
+            ],
+        ) -> Callable[Concatenate[Self, Args], ReturnType]: ...
         def with_state_decorator(
             func: Callable[
                 Concatenate[SelectorOutput, Args],
                 ReturnType,
+            ]
+            | Callable[
+                Concatenate[Self, SelectorOutput, Args],
+                ReturnType,
             ],
-        ) -> Callable[Args, ReturnType]:
+        ) -> Callable[Args, ReturnType] | Callable[Concatenate[Self, Args], ReturnType]:
+            signature = drop_with_store_parameter(func)
+
+            if (
+                signature.parameters
+                and next(iter(signature.parameters.values())).name == 'self'
+            ):
+                func_ = cast(
+                    'Callable[Concatenate[Self, SelectorOutput, Args], ReturnType]',
+                    func,
+                )
+
+                def wrapper(*args: Args.args, **kwargs: Args.kwargs) -> ReturnType:
+                    if self._state is None:
+                        msg = 'Store has not been initialized yet.'
+                        raise RuntimeError(msg)
+                    self_ = cast('Self', args[0])
+                    args_ = cast('Any', args[1:])
+                    return func_(self_, selector(self._state), *args_, **kwargs)
+
+                wrapped = wraps(func_)(wrapper)
+                wrapped.__signature__ = signature  # pyright: ignore [reportAttributeAccessIssue]
+
+                return wrapped
+
+            func_ = cast(
+                'Callable[Concatenate[SelectorOutput, Args], ReturnType]',
+                func,
+            )
+
             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)
+                return func_(selector(self._state), *args, **kwargs)
 
-            wrapped = wraps(func)(wrapper)
-            wrapped.__signature__ = drop_with_store_parameter(func)  # pyright: ignore [reportAttributeAccessIssue]
+            wrapped = wraps(func_)(wrapper)
+            wrapped.__signature__ = signature  # pyright: ignore [reportAttributeAccessIssue]
 
             return wrapped
 

tests/test_async.py

@@ -101,7 +101,7 @@ async def sync_mirror(value: int) -> int:
     async def _(mirrored_value: int) -> None:
         if mirrored_value < INCREMENTS:
             return
-        await asyncio.sleep(0.1)
+        await asyncio.sleep(0.001)
         store.dispatch(FinishAction())
 
     dispatch_actions(store)
@@ -124,7 +124,7 @@ async def _(values: tuple[int, int]) -> None:
         elif value < INCREMENTS:
             store.dispatch(IncrementAction())
         else:
-            await asyncio.sleep(0.1)
+            await asyncio.sleep(0.001)
             store.dispatch(FinishAction())
 
 
@@ -139,6 +139,7 @@ async def sync_mirror(value: int) -> int:
     @store.autorun(lambda state: (state.value, state.mirrored_value))
     async def _(values: tuple[int, int]) -> None:
         value, mirrored_value = values
+        await asyncio.sleep(0.001)
         if mirrored_value != value:
             sync_mirror_returned_value = sync_mirror()
             assert 'awaited=False' in str(sync_mirror_returned_value)
@@ -148,7 +149,6 @@ async def _(values: tuple[int, int]) -> None:
         elif value < INCREMENTS:
             store.dispatch(IncrementAction())
         else:
-            await asyncio.sleep(0.1)
             store.dispatch(FinishAction())
 
 
@@ -165,7 +165,7 @@ async def _(value: int) -> int:
     async def _(mirrored_value: int) -> None:
         if mirrored_value < INCREMENTS:
             return
-        await asyncio.sleep(0.1)
+        await asyncio.sleep(0.001)
         store.dispatch(FinishAction())
 
     dispatch_actions(store)
@@ -184,7 +184,7 @@ async def _(value: int) -> None:
         assert await doubled() == value * 2
         for _ in range(10):
             await doubled()
-        await asyncio.sleep(0.01)
+        await asyncio.sleep(0.001)
         if value < INCREMENTS:
             store.dispatch(IncrementAction())
         else:
@@ -209,7 +209,7 @@ async def _(value: int) -> None:
         calls_length = len(calls)
         assert await doubled() == value * 2
         assert len(calls) == calls_length + 1
-        await asyncio.sleep(0.01)
+        await asyncio.sleep(0.001)
 
         if value < INCREMENTS:
             store.dispatch(IncrementAction())
@@ -230,7 +230,7 @@ async def multiplied(value: int, factor: int) -> int:
     async def _(value: int) -> None:
         assert await multiplied(factor=2) == value * 2
         assert await multiplied(factor=3) == value * 3
-        await asyncio.sleep(0.01)
+        await asyncio.sleep(0.001)
         if value < INCREMENTS:
             store.dispatch(IncrementAction())
         else:
@@ -251,7 +251,7 @@ async def doubled(value: int) -> int:
     @store.autorun(lambda state: state.value)
     async def _(value: int) -> None:
         assert await doubled() == value * 2
-        await asyncio.sleep(0.01)
+        await asyncio.sleep(0.001)
         if value < INCREMENTS:
             store.dispatch(IncrementAction())
         else:
@@ -265,7 +265,7 @@ def test_subscription(store: StoreType) -> None:
     async def render(state: StateType) -> None:
         if state.value == INCREMENTS:
             unsubscribe()
-            await asyncio.sleep(0.1)
+            await asyncio.sleep(0.001)
             store.dispatch(FinishAction())
 
     unsubscribe = store._subscribe(render)  # noqa: SLF001
@@ -277,7 +277,7 @@ def test_event_subscription(store: StoreType) -> None:
     async def handler(event: IncrementEvent) -> None:
         if event.post_value == INCREMENTS:
             unsubscribe()
-            await asyncio.sleep(0.1)
+            await asyncio.sleep(0.001)
             store.dispatch(FinishAction())
 
     unsubscribe = store.subscribe_event(IncrementEvent, handler)

tests/test_with_state.py

@@ -198,3 +198,24 @@ def check(self: X, value: int) -> None:
     store.dispatch(FinishAction())
 
     check_spy.assert_called_once_with(0)
+
+
+def test_methods(store: StoreType) -> None:
+    """Test `with_state` decorator with methods."""
+
+    class SomeClass:
+        @store.with_state(lambda state: state.value)
+        def some_method(self, value: int) -> int:
+            return value
+
+    instance = SomeClass()
+
+    store.dispatch(InitAction())
+
+    assert instance.some_method() == 0
+
+    store.dispatch(_IncrementAction())
+
+    assert instance.some_method() == 1
+
+    store.dispatch(FinishAction())