From cd3c26f0d073ac4bd4357e8f4db2d90060a054d2 Mon Sep 17 00:00:00 2001
From: sescer <borya981@gmail.com>
Date: Tue, 28 Jul 2026 17:37:52 +0700
Subject: [PATCH] fix: allow registering bound methods as tasks

Skip ProcessPoolExecutor name rewriting for non-functions so
broker.register_task works with bound methods again (#436).
---
 taskiq/decor.py          | 30 ++++++++++++--------
 tests/abc/test_broker.py | 59 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+), 12 deletions(-)

diff --git a/taskiq/decor.py b/taskiq/decor.py
index 8f90541c..5080cb78 100644
--- a/taskiq/decor.py
+++ b/taskiq/decor.py
@@ -1,3 +1,4 @@
+import inspect
 import sys
 from collections.abc import Callable, Coroutine
 from copy import copy
@@ -75,18 +76,23 @@ def __init__(
         # it back to the module where it was defined.
         # This way ProcessPoolExecutor will be able to import
         # the function by it's name and verify its correctness.
-        new_name = f"{original_func.__name__}__taskiq_original"
-        self.original_func.__name__ = new_name
-        if hasattr(self.original_func, "__qualname__"):
-            original_qualname = self.original_func.__qualname__.rsplit(".")
-            original_qualname[-1] = new_name
-            new_qualname = ".".join(original_qualname)
-            self.original_func.__qualname__ = new_qualname
-        setattr(
-            sys.modules[original_func.__module__],
-            new_name,
-            original_func,
-        )
+        #
+        # Bound methods are skipped: their `__name__` is not
+        # assignable, and renaming via `__func__` would mutate
+        # the class-level function and break pickle lookup.
+        if inspect.isfunction(original_func):
+            new_name = f"{original_func.__name__}__taskiq_original"
+            self.original_func.__name__ = new_name
+            if hasattr(self.original_func, "__qualname__"):
+                original_qualname = self.original_func.__qualname__.rsplit(".")
+                original_qualname[-1] = new_name
+                new_qualname = ".".join(original_qualname)
+                self.original_func.__qualname__ = new_qualname
+            setattr(
+                sys.modules[original_func.__module__],
+                new_name,
+                original_func,
+            )
 
     # Docs for this method are omitted in order to help
     # your IDE resolve correct docs for it.
diff --git a/tests/abc/test_broker.py b/tests/abc/test_broker.py
index 636f9576..45f1b5d5 100644
--- a/tests/abc/test_broker.py
+++ b/tests/abc/test_broker.py
@@ -1,8 +1,11 @@
 from collections.abc import AsyncGenerator
+from concurrent.futures import ProcessPoolExecutor
 from copy import copy
+from types import MethodType
 
 import pytest
 
+from taskiq import InMemoryBroker
 from taskiq.abc.broker import AsyncBroker
 from taskiq.decor import AsyncTaskiqDecoratedTask
 from taskiq.events import TaskiqEvents
@@ -30,6 +33,15 @@ async def listen(self) -> AsyncGenerator[BrokerMessage, None]:  # type: ignore
         """
 
 
+_process_pool_broker = _TestBroker()
+
+
+@_process_pool_broker.task(task_name="process_pool_sync_add")
+def process_pool_sync_add(value: int) -> int:
+    """Module-level sync task used by ProcessPoolExecutor regression test."""
+    return value + 1
+
+
 def test_decorator_success() -> None:
     """Test that decoration without parameters works."""
     tbrok = _TestBroker()
@@ -82,6 +94,53 @@ async def test_task() -> None: ...
     assert test_task.labels == old_labels
 
 
+def test_register_task_accepts_bound_method() -> None:
+    """Bound methods must register without mutating method.__name__."""
+    broker = _TestBroker()
+
+    class Counter:
+        def increment(self, value: int) -> int:
+            return value + 1
+
+    instance = Counter()
+    task = broker.register_task(instance.increment)
+
+    assert isinstance(task, AsyncTaskiqDecoratedTask)
+    assert isinstance(task.original_func, MethodType)
+    assert task.original_func.__self__ is instance
+    assert task.task_name == f"{instance.increment.__module__}:increment"
+    assert broker.find_task(task.task_name) is task
+
+
+@pytest.mark.anyio
+async def test_bound_method_task_executes_with_self() -> None:
+    """Registered bound method must keep instance state across execution."""
+    broker = InMemoryBroker(await_inplace=True)
+
+    class Counter:
+        def __init__(self) -> None:
+            self.total = 0
+
+        def add(self, value: int) -> int:
+            self.total += value
+            return self.total
+
+    instance = Counter()
+    task = broker.register_task(instance.add, task_name="bound.add")
+
+    kicked = await task.kiq(5)
+    result = await kicked.wait_result()
+    assert result.return_value == 5
+    assert instance.total == 5
+
+
+def test_process_pool_runs_module_level_decorated_sync_function() -> None:
+    """ProcessPoolExecutor must be able to run renamed original sync functions."""
+    with ProcessPoolExecutor(max_workers=1) as executor:
+        future = executor.submit(process_pool_sync_add.original_func, 41)
+        assert future.result(timeout=5) == 42
+
+
 @pytest.mark.anyio
 @pytest.mark.parametrize(
     ("is_worker_process", "startup", "shutdown"),