diff --git a/matrix/__init__.py b/matrix/__init__.py index 07bee19..1845a21 100644 --- a/matrix/__init__.py +++ b/matrix/__init__.py @@ -15,6 +15,7 @@ from .help import HelpCommand from .checks import cooldown from .room import Room +from .space import Space from .message import Message from .extension import Extension @@ -28,6 +29,7 @@ "HelpCommand", "cooldown", "Room", + "Space", "Message", "Extension", ] diff --git a/matrix/bot.py b/matrix/bot.py index edf9a2b..8879568 100644 --- a/matrix/bot.py +++ b/matrix/bot.py @@ -7,7 +7,7 @@ from nio import AsyncClient, Event, MatrixRoom -from .room import Room +from .room import Room, make_room from .space import Space from .group import Group from .config import Config @@ -100,8 +100,7 @@ def get_room(self, room_id: str) -> Room | None: ``` """ if matrix_room := self.client.rooms.get(room_id): - room_cls = Space if matrix_room.room_type == "m.space" else Room - return room_cls(matrix_room=matrix_room, client=self.client) + return make_room(matrix_room, self.client) return None def get_rooms(self) -> list[Room]: @@ -123,8 +122,7 @@ def get_rooms(self) -> list[Room]: rooms = [] for matrix_room in self.client.rooms.values(): - room_cls = Space if matrix_room.room_type == "m.space" else Room - rooms.append(room_cls(matrix_room=matrix_room, client=self.client)) + rooms.append(make_room(matrix_room, self.client)) return rooms diff --git a/matrix/room.py b/matrix/room.py index c737070..1b18593 100644 --- a/matrix/room.py +++ b/matrix/room.py @@ -17,10 +17,22 @@ ) from matrix.types import File, Image, Audio, Video +_registry: dict[str, type["Room"]] = {} + + +def make_room(matrix_room: MatrixRoom, client: AsyncClient) -> "Room": + room_cls = _registry.get(matrix_room.room_type, Room) + return room_cls(matrix_room, client) + class Room: """Represents a Matrix room and provides methods to interact with it.""" + def __init_subclass__(cls, room_type: str | None = None, **kwargs: Any) -> None: + super().__init_subclass__(**kwargs) + if room_type: + _registry[room_type] = cls + def __init__(self, matrix_room: MatrixRoom, client: AsyncClient) -> None: self._matrix_room: MatrixRoom = matrix_room self._client: AsyncClient = client diff --git a/matrix/space.py b/matrix/space.py index b22d28a..4e20189 100644 --- a/matrix/space.py +++ b/matrix/space.py @@ -1,5 +1,27 @@ -from .room import Room +from matrix.room import Room, make_room -class Space(Room): - pass +class Space(Room, room_type="m.space"): + def get_children(self) -> list[Room]: + """Return the child rooms and spaces of this space. + + ## Example + + ```python + space = bot.get_space("!space123:matrix.org") + + for child in space.get_children(): + print(child.name) + ``` + """ + children = [] + + for room_id in self.children: + matrix_room = self._client.rooms.get(room_id) + + if not matrix_room: + continue + + children.append(make_room(matrix_room, self._client)) + + return children diff --git a/tests/test_bot.py b/tests/test_bot.py index bf0b8e6..2125ac0 100644 --- a/tests/test_bot.py +++ b/tests/test_bot.py @@ -3,7 +3,7 @@ from unittest.mock import AsyncMock, MagicMock, patch from nio import MatrixRoom, RoomMessageText -from matrix.bot import Bot, Config, Extension, Room, Space +from matrix import Bot, Config, Extension, Room, Space from matrix.errors import ( CheckError, CommandNotFoundError, diff --git a/tests/test_room.py b/tests/test_room.py index eafe95b..bd2cb4c 100644 --- a/tests/test_room.py +++ b/tests/test_room.py @@ -2,7 +2,8 @@ from unittest.mock import AsyncMock, Mock, MagicMock from nio import MatrixRoom, Event from matrix.errors import MatrixError -from matrix.room import Room +from matrix.room import Room, make_room +from matrix.space import Space from matrix.message import Message @@ -387,3 +388,35 @@ def test_room_client_property__expect_async_client(room, client): def test_room_unknown_attribute__expect_attribute_error(room): with pytest.raises(AttributeError): _ = room.nonexistent_attribute + + +def test_make_room__with_regular_room__expect_room_instance(client): + matrix_room = MatrixRoom( + room_id="!room:example.com", own_user_id="@bot:example.com" + ) + + result = make_room(matrix_room, client) + + assert type(result) is Room + + +def test_make_room__with_space_room__expect_space_instance(client): + matrix_room = MatrixRoom( + room_id="!space:example.com", own_user_id="@bot:example.com" + ) + matrix_room.room_type = "m.space" + + result = make_room(matrix_room, client) + + assert type(result) is Space + + +def test_make_room__with_unknown_room_type__expect_room_instance(client): + matrix_room = MatrixRoom( + room_id="!room:example.com", own_user_id="@bot:example.com" + ) + matrix_room.room_type = "m.unknown" + + result = make_room(matrix_room, client) + + assert type(result) is Room diff --git a/tests/test_space.py b/tests/test_space.py new file mode 100644 index 0000000..f171ff9 --- /dev/null +++ b/tests/test_space.py @@ -0,0 +1,94 @@ +import pytest +from unittest.mock import AsyncMock, Mock +from nio import MatrixRoom +from matrix.room import Room +from matrix.space import Space + + +@pytest.fixture +def client(): + return AsyncMock() + + +@pytest.fixture +def matrix_space(): + space = MatrixRoom(room_id="!space:example.com", own_user_id="@bot:example.com") + space.name = "Test Space" + space.room_type = "m.space" + return space + + +@pytest.fixture +def space(matrix_space, client): + return Space(matrix_space, client) + + +def test_get_children__with_room_child__expect_room_instance( + space, matrix_space, client +): + child = MatrixRoom(room_id="!child:example.com", own_user_id="@bot:example.com") + child.name = "Child Room" + matrix_space.children = {"!child:example.com"} + client.rooms = {"!child:example.com": child} + + result = space.get_children() + + assert len(result) == 1 + assert type(result[0]) is Room + assert result[0].room_id == "!child:example.com" + + +def test_get_children__with_space_child__expect_space_instance( + space, matrix_space, client +): + child = MatrixRoom(room_id="!subspace:example.com", own_user_id="@bot:example.com") + child.name = "Sub Space" + child.room_type = "m.space" + matrix_space.children = {"!subspace:example.com"} + client.rooms = {"!subspace:example.com": child} + + result = space.get_children() + + assert len(result) == 1 + assert isinstance(result[0], Space) + assert result[0].room_id == "!subspace:example.com" + + +def test_get_children__with_unjoined_child__expect_child_skipped( + space, matrix_space, client +): + matrix_space.children = {"!unknown:example.com"} + client.rooms = {} + + result = space.get_children() + + assert result == [] + + +def test_get_children__with_no_children__expect_empty_list(space, matrix_space, client): + matrix_space.children = set() + client.rooms = {} + + result = space.get_children() + + assert result == [] + + +def test_get_children__with_mixed_children__expect_correct_types( + space, matrix_space, client +): + room_child = MatrixRoom(room_id="!room:example.com", own_user_id="@bot:example.com") + space_child = MatrixRoom(room_id="!sub:example.com", own_user_id="@bot:example.com") + space_child.room_type = "m.space" + matrix_space.children = {"!room:example.com", "!sub:example.com"} + client.rooms = { + "!room:example.com": room_child, + "!sub:example.com": space_child, + } + + result = space.get_children() + + assert len(result) == 2 + types = {r.room_id: type(r) for r in result} + assert types["!room:example.com"] is Room + assert types["!sub:example.com"] is Space