Space: recursive get_children

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)
 
 

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,16 +17,29 @@ 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)
 
             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

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