Skip to content

Commit 9ec0aae

Browse files
authored
feat: Return request handlers from decorator methods to allow further decoration (#934)
### Description Currently, wrapping a handler with a decorator nullifies the function because it returns `None` [in the code](https://github.com/apify/crawlee-python/blob/28579996e454e19dd69ed6ee09e4293d4e056a6e/src/crawlee/router.py#L44-L45). See the "None" popup over here: ![image](https://github.com/user-attachments/assets/96df44c5-3ceb-4145-9bdd-f697bdf2da3d) By having the router.handler function return the original input handler, we can now reuse the handler function again in other places as it's not set to `None` anymore. This allows for invoking the original function even after applying the decorator and for adding multiple decorators to the same handler function. The latter for example allows for using the same handler function for multiple labels (a use case I ran into). ### Issues Not applicable. ### Testing - [x] Added a test to ensure handler is not nullified - [x] Added a test for multiple labels assigned to a single handler ### Checklist - [x] CI passed Please let me know your thoughts! 🙌
1 parent 2fdfd87 commit 9ec0aae

2 files changed

Lines changed: 31 additions & 2 deletions

File tree

src/crawlee/router.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,16 +33,19 @@ def default_handler(self: Router, handler: RequestHandler[TCrawlingContext]) ->
3333

3434
return handler
3535

36-
def handler(self, label: str) -> Callable[[RequestHandler[TCrawlingContext]], None]:
36+
def handler(
37+
self, label: str
38+
) -> Callable[[RequestHandler[TCrawlingContext]], Callable[[TCrawlingContext], Awaitable]]:
3739
"""A decorator used to register a label-based handler.
3840
3941
The registered will be invoked only for requests with the exact same label.
4042
"""
4143
if label in self._handlers_by_label:
4244
raise RuntimeError(f'A handler for label `{label}` is already registered')
4345

44-
def wrapper(handler: Callable[[TCrawlingContext], Awaitable]) -> None:
46+
def wrapper(handler: Callable[[TCrawlingContext], Awaitable]) -> Callable[[TCrawlingContext], Awaitable]:
4547
self._handlers_by_label[label] = handler
48+
return handler
4649

4750
return wrapper
4851

tests/unit/test_router.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,3 +89,29 @@ async def default_handler(_context: MockContext) -> None:
8989
mock_default_handler.assert_not_called()
9090
mock_handler_a.assert_not_called()
9191
mock_handler_b.assert_called()
92+
93+
94+
async def test_router_handler_not_nullified() -> None:
95+
router = Router[MockContext]()
96+
97+
@router.handler('A')
98+
async def handler_a(_context: MockContext) -> None:
99+
pass
100+
101+
assert handler_a is not None
102+
103+
104+
async def test_router_multi_labelled_handler() -> None:
105+
router = Router[MockContext]()
106+
mock_handler = Mock()
107+
108+
@router.handler('A')
109+
@router.handler('B')
110+
async def handler(_context: MockContext) -> None:
111+
mock_handler(_context.request.label)
112+
113+
await router(MockContext(label='A'))
114+
mock_handler.assert_called_with('A')
115+
await router(MockContext(label='B'))
116+
mock_handler.assert_called_with('B')
117+
assert mock_handler.call_count == 2

0 commit comments

Comments
 (0)