Add handle_response to normalize @on_method return values (#15)

Garulf · web-flow · commit 3e08c2e3857f · 2026-06-25T22:01:53.000-04:00
* Add handle_response to normalize @on_method return values

Plugin methods can now yield/return Result, List[Result], or use
generators without calling send_results() explicitly. Introduces
types.py to break the result.py→plugin.py circular import, and
response.py with handle_response() applied transparently in the
on_method wrapper. Async generators are handled in EventHandler.

* Fix missing Callable import in plugin.py after Method alias moved to types.py

* Update docs and examples to use yield-based result pattern
diff --git a/README.md b/README.md
@@ -28,25 +28,25 @@ python -m pip install pyflowlauncher[all]
 A basic plugin using a function as the query method.
 
 ```py
-from pyflowlauncher import Plugin, Result, send_results
-from .models.json_rpc import JsonRPCResponse
+from pyflowlauncher import Plugin, Result
 
 plugin = Plugin()
 
 
 @plugin.on_method
-def query(query: str) -> JsonRPCResponse:
-    r = Result(
+def query(query: str):
+    yield Result(
         title="This is a title!",
         subtitle="This is the subtitle!",
         icon="icon.png"
     )
-    return send_results([r])
 
 
 plugin.run()
 ```
 
+Methods decorated with `@plugin.on_method` can `yield` one or more `Result` objects, return a list of `Result` objects, or return a single `Result` — the framework normalizes all forms into the correct response automatically.
+
 ### Advanced plugin
 
 A more advanced usage using a `Method` class as the query method.
diff --git a/docs/examples/guide/plugin_methods/example1.py b/docs/examples/guide/plugin_methods/example1.py
@@ -1,17 +1,15 @@
-from pyflowlauncher import Plugin, Result, send_results
-from pyflowlauncher.models.json_rpc import JsonRPCResponse
+from pyflowlauncher import Plugin, Result
 
 plugin = Plugin()
 
 
 @plugin.on_method
-def query(query: str) -> JsonRPCResponse:
-    r = Result(
+def query(query: str):
+    yield Result(
         title="This is a title!",
         subtitle="This is the subtitle!",
         json_rpc_action={"Method": "action", "Parameters": []}
     )
-    return send_results([r])
 
 
 @plugin.on_method
diff --git a/docs/examples/guide/plugin_methods/example2.py b/docs/examples/guide/plugin_methods/example2.py
@@ -1,17 +1,16 @@
-from pyflowlauncher import Plugin, Result, send_results
-from pyflowlauncher.models.json_rpc import JsonRPCResponse
+from pyflowlauncher import Plugin, Result
 
 plugin = Plugin()
 
 
 @plugin.on_method
-def query(query: str) -> JsonRPCResponse:
+def query(query: str):
     r = Result(
         title="This is a title!",
         subtitle="This is the subtitle!",
     )
     r.add_action(action, ["stuff"])
-    return send_results([r])
+    yield r
 
 
 def action(params: list[str]):
diff --git a/docs/examples/guide/scoring_results/example1.py b/docs/examples/guide/scoring_results/example1.py
@@ -1,20 +1,19 @@
-from pyflowlauncher import Plugin, Result, send_results
-from pyflowlauncher.models.json_rpc import JsonRPCResponse
+from pyflowlauncher import Plugin, Result
 from pyflowlauncher.utils import score_results
 
 plugin = Plugin()
 
 
 @plugin.on_method
-def query(query: str) -> JsonRPCResponse:
-    results = []
-    for _ in range(100):
-        r = Result(
+def query(query: str):
+    results = [
+        Result(
             title="This is a title!",
             subtitle="This is the subtitle!",
         )
-        results.append(r)
-    return send_results(score_results(query, results))
+        for _ in range(100)
+    ]
+    yield from score_results(query, results)
 
 
 plugin.run()
diff --git a/docs/guide/Plugin methods.md b/docs/guide/Plugin methods.md
@@ -2,15 +2,42 @@
 
 Flow Launcher can call custom methods created by your plugin as well. To do so simply register the method with your plugin.
 
-You can register any Function by using the `on_method` decorator or by using the `add_method` method from `Plugin`.
+You can register any function by using the `on_method` decorator or by using the `add_method` method from `Plugin`.
+
+## Returning results
+
+Methods decorated with `@plugin.on_method` support several return styles — the framework normalizes all of them automatically:
+
+```py
+# yield a single result
+@plugin.on_method
+def query(query: str):
+    yield Result(title="Hello!")
+
+# yield multiple results
+@plugin.on_method
+def query(query: str):
+    for item in data:
+        yield Result(title=item)
+
+# return a list of results
+@plugin.on_method
+def query(query: str):
+    return [Result(title="a"), Result(title="b")]
+
+# return a single result
+@plugin.on_method
+def query(query: str):
+    return Result(title="Hello!")
+```
 
 ## Example 1
 
 ```py
 --8<-- "docs/examples/guide/plugin_methods/example1.py"
 ```
 
-Alternativley you can register and add the Method to a Result in one line by using `action` method.
+Alternatively you can register and add the Method to a Result in one line by using the `action` method.
 
 ## Example 2
 
diff --git a/pyflowlauncher/__init__.py b/pyflowlauncher/__init__.py
@@ -3,6 +3,7 @@
 from .plugin import Plugin
 from .result import Result, send_results
 from .method import Method
+from .response import handle_response
 
 
 logger = logging.getLogger(__name__)
@@ -13,4 +14,5 @@
     "send_results",
     "Result",
     "Method",
+    "handle_response",
 ]
diff --git a/pyflowlauncher/event.py b/pyflowlauncher/event.py
@@ -1,6 +1,9 @@
 import asyncio
+import inspect
 from typing import Any, Callable, Iterable, Type, Union
 
+from .result import Result, send_results
+
 
 class EventNotFound(Exception):
 
@@ -39,6 +42,14 @@ def get_event(self, event: str) -> Callable[..., Any]:
     async def _await_maybe(self, result: Any) -> Any:
         if asyncio.iscoroutine(result):
             return await result
+        if inspect.isasyncgen(result):
+            results = []
+            async for item in result:
+                if isinstance(item, Result):
+                    results.append(item)
+                elif isinstance(item, list):
+                    results.extend(r for r in item if isinstance(r, Result))
+            return send_results(results)
         return result
 
     async def trigger_exception_handler(self, exception: Exception) -> Any:
diff --git a/pyflowlauncher/plugin.py b/pyflowlauncher/plugin.py
@@ -9,12 +9,13 @@
 from .base import pyFlowLauncherObject
 
 from .event import EventHandler
+from .response import handle_response
 from .jsonrpc import JsonRPCClient, JsonRPCRequest
 from .models.plugin_manifest import FILE_NAME
 from .manifest import Manifest
 
 
-Method = Callable[..., Any]
+from .types import Method
 
 
 class Plugin(pyFlowLauncherObject):
@@ -38,7 +39,7 @@ def add_methods(self, methods: Iterable[Method]) -> None:
     def on_method(self, method: Method) -> Method:
         @wraps(method)
         def wrapper(*args, **kwargs):
-            return method(*args, **kwargs)
+            return handle_response(method(*args, **kwargs))
         self.add_method(wrapper)
         return wrapper
 
diff --git a/pyflowlauncher/response.py b/pyflowlauncher/response.py
@@ -0,0 +1,33 @@
+import inspect
+from typing import Any, Generator, Union
+
+from .result import Result, send_results
+from .models.json_rpc import JsonRPCRequest, JsonRPCResponse
+
+
+def handle_response(result: Any) -> Union[JsonRPCResponse, JsonRPCRequest, None]:
+    """Normalize a method's return value into a JSON-RPC response.
+
+    Accepts: Result, list of Result, generator of Result/list, JsonRPCRequest,
+    JsonRPCResponse, coroutine, async generator, or None.
+    Coroutines and async generators are returned as-is for the event loop to handle.
+    """
+    if result is None:
+        return None
+    if isinstance(result, Result):
+        return send_results([result])
+    if isinstance(result, list):
+        return send_results(result)
+    if inspect.isgenerator(result):
+        return _collect_generator(result)
+    return result
+
+
+def _collect_generator(gen: Generator) -> JsonRPCResponse:
+    results = []
+    for item in gen:
+        if isinstance(item, Result):
+            results.append(item)
+        elif isinstance(item, list):
+            results.extend(r for r in item if isinstance(r, Result))
+    return send_results(results)
diff --git a/pyflowlauncher/result.py b/pyflowlauncher/result.py
@@ -4,7 +4,7 @@
 from pathlib import Path
 from typing import Any, Callable, Dict, Iterable, List, Optional, Union, cast
 
-from .plugin import Method
+from .types import Method
 from .models.result import Glyph, PreviewInfo
 from .models.json_rpc import JsonRPCResult, JsonRPCRequest, JsonRPCResponse
 
diff --git a/pyflowlauncher/types.py b/pyflowlauncher/types.py
@@ -0,0 +1,3 @@
+from typing import Any, Callable
+
+Method = Callable[..., Any]
diff --git a/tests/test_response.py b/tests/test_response.py
@@ -0,0 +1,139 @@
+import pytest
+from pyflowlauncher import Plugin, Result, handle_response
+from pyflowlauncher.result import send_results
+
+
+# --- Unit tests for handle_response ---
+
+def test_none():
+    assert handle_response(None) is None
+
+
+def test_single_result():
+    r = Result(title="hello")
+    assert handle_response(r) == send_results([r])
+
+
+def test_list_of_results():
+    results = [Result(title="a"), Result(title="b")]
+    assert handle_response(results) == send_results(results)
+
+
+def test_generator_single_result():
+    def gen():
+        yield Result(title="a")
+
+    assert handle_response(gen()) == send_results([Result(title="a")])
+
+
+def test_generator_multiple_results():
+    def gen():
+        yield Result(title="a")
+        yield Result(title="b")
+
+    assert handle_response(gen()) == send_results([Result(title="a"), Result(title="b")])
+
+
+def test_generator_yields_list():
+    def gen():
+        yield [Result(title="a"), Result(title="b")]
+
+    assert handle_response(gen()) == send_results([Result(title="a"), Result(title="b")])
+
+
+def test_generator_mixed_yields():
+    def gen():
+        yield Result(title="a")
+        yield [Result(title="b"), Result(title="c")]
+
+    assert handle_response(gen()) == send_results([
+        Result(title="a"), Result(title="b"), Result(title="c")
+    ])
+
+
+def test_jsonrpc_request_passthrough():
+    req = {"Method": "Flow.Launcher.HideApp", "Parameters": []}
+    assert handle_response(req) is req
+
+
+def test_jsonrpc_response_passthrough():
+    resp = send_results([Result(title="x")])
+    assert handle_response(resp) is resp
+
+
+# --- Integration: @plugin.on_method with generator ---
+
+def test_on_method_generator_integration():
+    plugin = Plugin()
+
+    @plugin.on_method
+    def query(q: str):
+        yield Result(title="foo")
+        yield Result(title="bar")
+
+    import asyncio
+    result = asyncio.run(plugin._event_handler.trigger_event("query", "test"))
+    assert result == send_results([Result(title="foo"), Result(title="bar")])
+
+
+def test_on_method_single_result_integration():
+    plugin = Plugin()
+
+    @plugin.on_method
+    def query(q: str):
+        return Result(title="only")
+
+    import asyncio
+    result = asyncio.run(plugin._event_handler.trigger_event("query", "test"))
+    assert result == send_results([Result(title="only")])
+
+
+def test_on_method_list_integration():
+    plugin = Plugin()
+
+    @plugin.on_method
+    def query(q: str):
+        return [Result(title="a"), Result(title="b")]
+
+    import asyncio
+    result = asyncio.run(plugin._event_handler.trigger_event("query", "test"))
+    assert result == send_results([Result(title="a"), Result(title="b")])
+
+
+# --- Async generator support ---
+
+@pytest.mark.asyncio
+async def test_async_generator_single():
+    plugin = Plugin()
+
+    @plugin.on_method
+    async def query(q: str):
+        yield Result(title="async-foo")
+
+    result = await plugin._event_handler.trigger_event("query", "test")
+    assert result == send_results([Result(title="async-foo")])
+
+
+@pytest.mark.asyncio
+async def test_async_generator_multiple():
+    plugin = Plugin()
+
+    @plugin.on_method
+    async def query(q: str):
+        yield Result(title="a")
+        yield Result(title="b")
+
+    result = await plugin._event_handler.trigger_event("query", "test")
+    assert result == send_results([Result(title="a"), Result(title="b")])
+
+
+@pytest.mark.asyncio
+async def test_async_generator_yields_list():
+    plugin = Plugin()
+
+    @plugin.on_method
+    async def query(q: str):
+        yield [Result(title="x"), Result(title="y")]
+
+    result = await plugin._event_handler.trigger_event("query", "test")
+    assert result == send_results([Result(title="x"), Result(title="y")])
