allow strings or functions for tab filter function in python module

browser_cli/__init__.py

@@ -16,6 +16,8 @@ Usage:
     # When multiple browser instances are active, pass the alias:
     b = BrowserCLI(browser="brave")
 """
+from collections.abc import Callable, Iterable
+
 from browser_cli.client import BrowserNotConnected, send_command
 from browser_cli.models import Group, Tab
 
@@ -119,9 +121,11 @@ class BrowserCLI:
         """Switch browser focus to a tab by ID."""
         self._cmd("tabs.active", {"tabId": tab_id})
 
-    def tabs_filter(self, pattern: str) -> list[Tab]:
-        """Return tabs whose URL contains *pattern*."""
-        return [self._make_tab(t) for t in (self._cmd("tabs.filter", {"pattern": pattern}) or [])]
+    def tabs_filter(self, pattern_or_filter: str | Callable[[Tab], bool] | Callable[[list[Tab]], Iterable[Tab]]) -> list[Tab]:
+        """Return tabs filtered by pattern or a Python callable."""
+        if isinstance(pattern_or_filter, str):
+            return [self._make_tab(t) for t in (self._cmd("tabs.filter", {"pattern": pattern_or_filter}) or [])]
+        return self._apply_tab_filter(pattern_or_filter)
 
     def tabs_count(self, pattern: str | None = None) -> int:
         """Count open tabs, optionally filtered by URL pattern."""
@@ -267,3 +271,27 @@ class BrowserCLI:
 
     def clients(self) -> list[dict]:
         return self._cmd("clients.list", {})
+
+    def _apply_tab_filter(self, filter_fn: Callable[[Tab], bool] | Callable[[list[Tab]], Iterable[Tab]]) -> list[Tab]:
+        tabs = self.tabs_list()
+
+        try:
+            transformed = filter_fn(tabs)
+        except Exception:
+            transformed = None
+
+        if isinstance(transformed, list):
+            return transformed
+        if isinstance(transformed, tuple):
+            return list(transformed)
+        if isinstance(transformed, set):
+            return list(transformed)
+        if transformed is tabs:
+            return tabs
+        if isinstance(transformed, bool):
+            return [tab for tab in tabs if filter_fn(tab)]
+
+        try:
+            return list(transformed)
+        except TypeError:
+            return [tab for tab in tabs if filter_fn(tab)]

tests/test_api.py

@@ -238,6 +238,17 @@ class TestTabs:
         mock_send.return_value = None
         assert b.tabs_filter("x") == []
 
+    def test_tabs_filter_predicate(self, b, mock_send):
+        mock_send.return_value = [TAB_DATA, {**TAB_DATA, "id": 11, "url": "https://youtube.com"}]
+        tabs = b.tabs_filter(lambda tab: "youtube" in tab.url)
+        print(tabs)
+        assert [tab.id for tab in tabs] == [11]
+
+    def test_tabs_filter_list_transformer(self, b, mock_send):
+        mock_send.return_value = [TAB_DATA, {**TAB_DATA, "id": 11, "url": "https://example.com"}]
+        tabs = b.tabs_filter(lambda tabs: tabs[:1])
+        assert [tab.id for tab in tabs] == [10]
+
     def test_tabs_count(self, b, mock_send):
         mock_send.return_value = 5
         assert b.tabs_count() == 5