Merge pull request #353 from Tishka17/feature/or_widgets

support | for kbd and media

src/aiogram_dialog/api/entities/media.py

@@ -36,3 +36,15 @@ def __init__(
         self.file_id = file_id
         self.use_pipe = use_pipe
         self.kwargs = kwargs
+
+    def __eq__(self, other):
+        if type(other) is not type(self):
+            return False
+        return (
+            self.type == other.type and
+            self.url == other.url and
+            self.path == other.path and
+            self.file_id == other.file_id and
+            self.use_pipe == other.use_pipe and
+            self.kwargs == other.kwargs
+        )

src/aiogram_dialog/widgets/kbd/base.py

@@ -109,3 +109,46 @@ async def _process_other_callback(
         Can be used for layouts
         """
         return False
+
+    def __or__(self, other: "Keyboard") -> "Or":
+        # reduce nesting
+        if isinstance(other, Or):
+            return NotImplemented
+        return Or(self, other)
+
+    def __ror__(self, other: "Keyboard") -> "Or":
+        # reduce nesting
+        return Or(other, self)
+
+
+class Or(Keyboard):
+    def __init__(self, *widgets: Keyboard):
+        super().__init__()
+        self.widgets = widgets
+
+    async def _render_keyboard(
+            self, data: dict, manager: DialogManager,
+    ) -> RawKeyboard:
+        for widget in self.widgets:
+            res = await widget.render_keyboard(data, manager)
+            if res and any(res):
+                return res
+        return []
+
+    def __ior__(self, other: Keyboard) -> "Or":
+        self.widgets += (other,)
+        return self
+
+    def __or__(self, other: Keyboard) -> "Or":
+        # reduce nesting
+        return Or(*self.widgets, other)
+
+    def __ror__(self, other: Keyboard) -> "Or":
+        # reduce nesting
+        return Or(other, *self.widgets)
+
+    def find(self, widget_id: str) -> Optional[Keyboard]:
+        for text in self.widgets:
+            if found := text.find(widget_id):
+                return found
+        return None

src/aiogram_dialog/widgets/media/base.py

@@ -21,3 +21,40 @@ async def _render_media(
             self, data: dict, manager: DialogManager,
     ) -> Optional[MediaAttachment]:
         return None
+
+    def __or__(self, other: "Media") -> "Or":
+        # reduce nesting
+        if isinstance(other, Or):
+            return NotImplemented
+        return Or(self, other)
+
+    def __ror__(self, other: "Media") -> "Or":
+        # reduce nesting
+        return Or(other, self)
+
+
+class Or(Media):
+    def __init__(self, *widgets: Media):
+        super().__init__()
+        self.widgets = widgets
+
+    async def _render_media(
+            self, data: dict, manager: DialogManager,
+    ) -> Optional[MediaAttachment]:
+        for widget in self.widgets:
+            res = await widget.render_media(data, manager)
+            if res:
+                return res
+        return None
+
+    def __ior__(self, other: Media) -> "Or":
+        self.widgets += (other,)
+        return self
+
+    def __or__(self, other: Media) -> "Or":
+        # reduce nesting
+        return Or(*self.widgets, other)
+
+    def __ror__(self, other: Media) -> "Or":
+        # reduce nesting
+        return Or(other, *self.widgets)

src/aiogram_dialog/widgets/text/base.py

@@ -37,6 +37,8 @@ async def _render_text(self, data, manager: DialogManager) -> str:
     def __add__(self, other: Union["Text", str]):
         if isinstance(other, str):
             other = Const(other)
+        elif isinstance(other, Multi):
+            return NotImplemented
         return Multi(self, other, sep="")
 
     def __radd__(self, other: Union["Text", str]):
@@ -47,6 +49,8 @@ def __radd__(self, other: Union["Text", str]):
     def __or__(self, other: Union["Text", str]):
         if isinstance(other, str):
             other = Const(other)
+        elif isinstance(other, Or):
+            return NotImplemented
         return Or(self, other)
 
     def __ror__(self, other: Union["Text", str]):
@@ -141,7 +145,7 @@ def __or__(self, other: Union[Text, str]) -> "Or":
 
     def __ror__(self, other: Union[Text, str]) -> "Or":
         if isinstance(other, str):
-            return Const(other)
+            other = Const(other)
         # reduce nesting
         return Or(other, *self.texts)
 

tests/widgets/kbd/test_base.py

@@ -0,0 +1,62 @@
+from unittest.mock import Mock
+
+import pytest
+from aiogram import F
+from aiogram.fsm.state import State
+from aiogram.types import KeyboardButton
+
+from aiogram_dialog import DialogManager
+from aiogram_dialog.api.entities import Context
+from aiogram_dialog.api.internal import RawKeyboard
+from aiogram_dialog.widgets.common import WhenCondition
+from aiogram_dialog.widgets.kbd import Keyboard
+
+
+@pytest.fixture()
+def mock_manager() -> DialogManager:
+    manager = Mock()
+    context = Context(
+        dialog_data={},
+        start_data={},
+        widget_data={},
+        state=State(),
+        _stack_id="_stack_id",
+        _intent_id="_intent_id",
+    )
+    manager.current_context = Mock(side_effect=lambda: context)
+
+    return manager
+
+
+class Button(Keyboard):
+    def __init__(self, id: str, when: WhenCondition = None):
+        super().__init__(when=when, id=id)
+
+    async def _render_keyboard(
+            self,
+            data,
+            manager: DialogManager,
+    ) -> RawKeyboard:
+        return [[KeyboardButton(text=self.widget_id)]]
+
+
+@pytest.mark.asyncio
+async def test_or(mock_manager):
+    text = Button("a") | Button("b")
+    res = await text.render_keyboard({}, mock_manager)
+    assert res == [[KeyboardButton(text="a")]]
+
+
+@pytest.mark.asyncio
+async def test_or_condition(mock_manager):
+    text = (
+        Button("A", when=F["a"]) |
+        Button("B", when=F["b"]) |
+        Button("C")
+    )
+    res = await text.render_keyboard({"a": True}, mock_manager)
+    assert res == [[KeyboardButton(text="A")]]
+    res = await text.render_keyboard({"b": True}, mock_manager)
+    assert res == [[KeyboardButton(text="B")]]
+    res = await text.render_keyboard({}, mock_manager)
+    assert res == [[KeyboardButton(text="C")]]

tests/widgets/media/test_base.py

@@ -0,0 +1,62 @@
+from unittest.mock import Mock
+
+import pytest
+from aiogram import F
+from aiogram.enums import ContentType
+from aiogram.fsm.state import State
+
+from aiogram_dialog import DialogManager
+from aiogram_dialog.api.entities import Context, MediaAttachment
+from aiogram_dialog.widgets.common import WhenCondition
+from aiogram_dialog.widgets.media import Media
+
+
+@pytest.fixture()
+def mock_manager() -> DialogManager:
+    manager = Mock()
+    context = Context(
+        dialog_data={},
+        start_data={},
+        widget_data={},
+        state=State(),
+        _stack_id="_stack_id",
+        _intent_id="_intent_id",
+    )
+    manager.current_context = Mock(side_effect=lambda: context)
+
+    return manager
+
+
+class Static(Media):
+    def __init__(self, path: str, when: WhenCondition = None):
+        super().__init__(when=when)
+        self.path = path
+
+    async def _render_media(
+            self,
+            data,
+            manager: DialogManager,
+    ) -> MediaAttachment:
+        return MediaAttachment(ContentType.PHOTO, path=self.path)
+
+
+@pytest.mark.asyncio
+async def test_or(mock_manager):
+    text = Static("a") | Static("b")
+    res = await text.render_media({}, mock_manager)
+    assert res == MediaAttachment(ContentType.PHOTO, path="a")
+
+
+@pytest.mark.asyncio
+async def test_or_condition(mock_manager):
+    text = (
+        Static("A", when=F["a"]) |
+        Static("B", when=F["b"]) |
+        Static("C")
+    )
+    res = await text.render_media({"a": True}, mock_manager)
+    assert res == MediaAttachment(ContentType.PHOTO, path="A")
+    res = await text.render_media({"b": True}, mock_manager)
+    assert res == MediaAttachment(ContentType.PHOTO, path="B")
+    res = await text.render_media({}, mock_manager)
+    assert res == MediaAttachment(ContentType.PHOTO, path="C")