diff --git a/matrix/room.py b/matrix/room.py index 10dda78..8aaa9d8 100644 --- a/matrix/room.py +++ b/matrix/room.py @@ -21,7 +21,7 @@ def make_room(matrix_room: MatrixRoom, client: AsyncClient) -> "Room": - room_cls = _registry.get(matrix_room.room_type, Room) + room_cls = _registry.get(str(matrix_room.room_type), Room) return room_cls(matrix_room, client) diff --git a/matrix/space.py b/matrix/space.py index c182902..525ee6f 100644 --- a/matrix/space.py +++ b/matrix/space.py @@ -1,11 +1,14 @@ +from typing import Self from matrix.room import Room, make_room class Space(Room, room_type="m.space"): - def get_children(self) -> list[Room]: + def get_children(self, depth: int = 1) -> list[Room | Self]: """Return the child rooms and spaces of this space that the bot has joined. - Children the bot has not joined are silently omitted. + Children the bot has not joined are silently omitted. Use `depth` to + recursively collect children of sub-spaces. `depth=1` returns direct + children only (default). ## Example @@ -14,9 +17,18 @@ def get_children(self) -> list[Room]: for child in space.get_children(): print(child.name) + + for child in space.get_children(depth=3): + print(child.name) ``` """ - children = [] + children: list[Room | Self] = [] + + if depth < 0: + raise ValueError(f"depth must be a non-negative integer, got {depth}") + + if depth == 0: + return [] for room_id in self.children: matrix_room = self._client.rooms.get(room_id) @@ -24,6 +36,10 @@ def get_children(self) -> list[Room]: if not matrix_room: continue - children.append(make_room(matrix_room, self._client)) + child = make_room(matrix_room, self._client) + children.append(child) + + if isinstance(child, Space) and depth > 1: + children.extend(child.get_children(depth - 1)) return children diff --git a/tests/test_space.py b/tests/test_space.py index f171ff9..444b6df 100644 --- a/tests/test_space.py +++ b/tests/test_space.py @@ -92,3 +92,60 @@ def test_get_children__with_mixed_children__expect_correct_types( types = {r.room_id: type(r) for r in result} assert types["!room:example.com"] is Room assert types["!sub:example.com"] is Space + + +def test_get_children__with_negative_depth__expect_value_error(space, matrix_space): + with pytest.raises(ValueError, match="depth must be a non-negative integer"): + space.get_children(depth=-1) + + +def test_get_children__with_depth_zero__expect_empty_list(space, matrix_space, client): + child = MatrixRoom(room_id="!child:example.com", own_user_id="@bot:example.com") + matrix_space.children = {"!child:example.com"} + client.rooms = {"!child:example.com": child} + + result = space.get_children(depth=0) + + assert result == [] + + +def test_get_children__with_depth_one__expect_no_recursion(space, matrix_space, client): + subspace = MatrixRoom( + room_id="!subspace:example.com", own_user_id="@bot:example.com" + ) + subspace.room_type = "m.space" + nested = MatrixRoom(room_id="!nested:example.com", own_user_id="@bot:example.com") + subspace.children = {"!nested:example.com"} + matrix_space.children = {"!subspace:example.com"} + client.rooms = { + "!subspace:example.com": subspace, + "!nested:example.com": nested, + } + + result = space.get_children(depth=1) + + assert len(result) == 1 + assert result[0].room_id == "!subspace:example.com" + + +def test_get_children__with_depth_two__expect_recursive_children( + space, matrix_space, client +): + subspace = MatrixRoom( + room_id="!subspace:example.com", own_user_id="@bot:example.com" + ) + subspace.room_type = "m.space" + nested = MatrixRoom(room_id="!nested:example.com", own_user_id="@bot:example.com") + subspace.children = {"!nested:example.com"} + matrix_space.children = {"!subspace:example.com"} + client.rooms = { + "!subspace:example.com": subspace, + "!nested:example.com": nested, + } + + result = space.get_children(depth=2) + + assert len(result) == 2 + ids = [r.room_id for r in result] + assert "!subspace:example.com" in ids + assert "!nested:example.com" in ids