Preserve query membership while a tag or component relation remains

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Removing a relation tag or component now preserves query membership when the other relation still exists.
 - 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/entity.py

@@ -9,6 +9,7 @@
     Final,
     Generic,
     TypeVar,
+    cast,
     overload,
 )
 from weakref import WeakKeyDictionary, WeakValueDictionary
@@ -700,6 +701,12 @@ def _relations_lookup_add(registry: Registry, origin: Entity, tag: object, targe
 
 def _relations_lookup_discard(registry: Registry, origin: Entity, tag: object, target: Entity) -> None:
     """Discard a relation tag/component from the lookup table and handle side effects."""
+    tags = registry._relation_tags_by_entity.get(origin)
+    if tags is not None and target in tags.get(tag, ()):
+        return  # A tag relation still owns this index entry.
+    components = registry._relation_components_by_entity.get(origin)
+    if components is not None and target in components.get(cast("ComponentKey[object]", tag), ()):
+        return  # A component relation still owns this index entry.
     registry._relations_lookup[(tag, target)].discard(origin)
     if not registry._relations_lookup[(tag, target)]:
         del registry._relations_lookup[(tag, target)]

tests/test_relations.py

@@ -2,12 +2,15 @@
 
 from __future__ import annotations
 
-from typing import Final
+from typing import TYPE_CHECKING, Final
 
 import pytest
 
 import tcod.ecs
 
+if TYPE_CHECKING:
+    from tcod.ecs.typing import ComponentKey
+
 # ruff: noqa: D103
 
 ChildOf: Final = "ChildOf"
@@ -326,3 +329,33 @@ def test_clear_with_relation() -> None:
     assert not child.components.keys()
     assert not child.tags
     assert not child.relation_components.keys()
+
+
+@pytest.mark.parametrize("key", [int, ("faction", int)])
+@pytest.mark.parametrize("remove_first", ["tag", "component"])
+def test_overlapping_relation_queries(key: ComponentKey[int], remove_first: str) -> None:
+    component_value = 0
+    world = tcod.ecs.Registry()
+    actor, faction = world["actor"], world["faction"]
+    actor.relation_tags_many[key].add(faction)
+    actor.relation_components[key][faction] = component_value
+    queries = [
+        world.Q.all_of(relations=[(key, faction)]),
+        world.Q.all_of(relations=[(key, ...)]),
+        world.Q.all_of(relations=[(actor, key, None)]),
+        world.Q.all_of(relations=[(..., key, None)]),
+    ]
+    expected = [{actor}, {actor}, {faction}, {faction}]
+    assert [set(q) for q in queries] == expected
+    if remove_first == "tag":
+        actor.relation_tags_many[key].remove(faction)
+        assert actor.relation_components[key][faction] == component_value
+    else:
+        del actor.relation_components[key][faction]
+        assert faction in actor.relation_tags_many[key]
+    assert [set(q) for q in queries] == expected
+    if remove_first == "tag":
+        del actor.relation_components[key][faction]
+    else:
+        actor.relation_tags_many[key].remove(faction)
+    assert not any(queries)