Resolve inherited components when collecting query rows

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Queries collecting component rows now resolve inherited components instead of raising `KeyError` for prefab instances.
 - Relation components can now be copied between distinct keys in the same entity, instead of the assignment being silently ignored.
   https://github.com/HexDecimal/python-tcod-ecs/pull/59
 

tcod/ecs/query.py

@@ -465,13 +465,16 @@ def __getitem__(self, key: tuple[ComponentKey[object], ...]) -> Iterable[tuple[A
         Entity = tcod.ecs.entity.Entity  # noqa: N806
 
         entities = list(self.all_of(components=set(key) - {Entity}).get_entities())
-        entity_components = []
+        entity_components: list[list[Any]] = []
         for component_key in key:
             if component_key is Entity:
                 entity_components.append(entities)
                 continue
             registry_components = self.registry._components_by_type[component_key]
-            entity_components.append([registry_components[entity] for entity in entities])
+            try:
+                entity_components.append([registry_components[entity] for entity in entities])
+            except KeyError:  # Some matching entities inherit this component.
+                entity_components.append([entity.components[component_key] for entity in entities])
         return zip(*entity_components, strict=True)
 
 

tests/test_traversal.py

@@ -2,15 +2,47 @@
 
 from __future__ import annotations
 
-from typing import Final
+from typing import TYPE_CHECKING, Final
 
 import pytest
 
-from tcod.ecs import IsA, Registry
+from tcod.ecs import Entity, IsA, Registry
+
+if TYPE_CHECKING:
+    from tcod.ecs.typing import ComponentKey
 
 # ruff: noqa: D103
 
 
+@pytest.mark.parametrize("key", [int, ("hp", int)])
+def test_query_component_traversal(key: ComponentKey[int]) -> None:
+    registry = Registry()
+    base = registry["base"]
+    base.components[key] = 2
+    child = base.instantiate()
+    grandchild = child.instantiate()
+    query = registry.Q.all_of(components=[key])
+    assert dict(query[Entity, key]) == {base: 2, child: 2, grandchild: 2}
+    assert list(registry.Q.all_of(components=[key], depth=0)[key,]) == [(2,)]
+
+    base.components[key] = 3
+    child.components[key] = 0
+    assert dict(query[Entity, key]) == {base: 3, child: 0, grandchild: 0}
+    assert set(query[key, key]) == {(3, 3), (0, 0)}
+    del child.components[key]
+    assert dict(query[Entity, key]) == {base: 3, child: 3, grandchild: 3}
+    del base.components[key]
+    assert list(query[key,]) == []
+
+
+def test_query_inherited_none() -> None:
+    registry = Registry()
+    base = registry["base"]
+    base.components[object] = None
+    child = base.instantiate()
+    assert dict(registry.Q[Entity, object]) == {base: None, child: None}
+
+
 def test_component_traversal() -> None:
     world = Registry()
     assert not world.Q.all_of(components=[str]).get_entities()