diff --git a/CHANGELOG.md b/CHANGELOG.md index aa44ea5..2d002e6 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/tcod/ecs/query.py b/tcod/ecs/query.py index cd4be36..5c42921 100644 --- a/tcod/ecs/query.py +++ b/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) diff --git a/tests/test_traversal.py b/tests/test_traversal.py index 950a98f..0bc8da3 100644 --- a/tests/test_traversal.py +++ b/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()