From 6fd355aa5fd6368c251b27ff119b9c6aec02298e Mon Sep 17 00:00:00 2001 From: Christopher Sumnicht Date: Tue, 18 Aug 2026 19:58:05 -0700 Subject: [PATCH] Added and fixed tests. --- src/punyecs/__init__.py | 67 ++++--- tests/test_basic.py | 2 +- tests/test_constraints.py | 161 +++++++++++++++++ tests/test_serialization.py | 56 ++++++ tests/test_world.py | 339 ++++++++++++++++++++++++++++++++++++ 5 files changed, 604 insertions(+), 21 deletions(-) create mode 100644 tests/test_constraints.py create mode 100644 tests/test_serialization.py create mode 100644 tests/test_world.py diff --git a/src/punyecs/__init__.py b/src/punyecs/__init__.py index ae00698..0d43957 100644 --- a/src/punyecs/__init__.py +++ b/src/punyecs/__init__.py @@ -1,7 +1,8 @@ from collections import namedtuple -from dataclasses import dataclass, field +from dataclasses import dataclass, field, asdict import json import operator +import sys from typing import Any, Callable def register_unary_op(op: Callable[[Any], Any]): @@ -26,17 +27,27 @@ def not_(self: Any): """Cannot override ``not``, so ``not_`` is used to make negations of Contraints. """ + if not hasattr(self, "eval"): + self = Const(self) return Constraint(unary_op=lambda ob: not ob, val1=self) -def and_(self: Any): +def and_(self: Any, other: Any): """Cannot override ``and``, so ``and_`` is used to conjunct two Contraints. """ - return Constraint(bin_op=lambda o, n: o and n, val1=self) + if not hasattr(self, "eval"): + self = Const(self) + if not hasattr(other, "eval"): + other = Const(other) + return Constraint(bin_op=lambda o, n: o and n, val1=self, val2=other) -def or_(self: Any): +def or_(self: Any, other: Any): """Cannot override ``or``, so ``or_`` is used to disjunct two Contraints.""" - return Constraint(bin_op=lambda o, n: o or n, val1=self) + if not hasattr(self, "eval"): + self = Const(self) + if not hasattr(other, "eval"): + other = Const(other) + return Constraint(bin_op=lambda o, n: o or n, val1=self, val2=other) @dataclass class c: @@ -97,15 +108,15 @@ def eval(self): __abs__ = register_unary_op(operator.__abs__) __add__ = register_bin_op(operator.__add__) - __radd__ = register_bin_op(operator.__add__) + __radd__ = register_bin_op(lambda o, n: operator.__add__(n, o)) __sub__ = register_bin_op(operator.__sub__) - __rsub__ = register_bin_op(operator.__sub__) + __rsub__ = register_bin_op(lambda o, n: operator.__sub__(n, o)) __mul__ = register_bin_op(operator.__mul__) - __rmul__ = register_bin_op(operator.__mul__) + __rmul__ = register_bin_op(lambda o, n: operator.__mul__(n, o)) __truediv__ = register_bin_op(operator.__truediv__) - __rtruediv__ = register_bin_op(operator.__truediv__) + __rtruediv__ = register_bin_op(lambda o, n: operator.__truediv__(n, o)) __floordiv__ = register_bin_op(operator.__floordiv__) - __rfloordiv__ = register_bin_op(operator.__floordiv__) + __rfloordiv__ = register_bin_op(lambda o, n: operator.__floordiv__(n, o)) # pyrefly: ignore __eq__ = register_bin_op(operator.__eq__) __lt__ = register_bin_op(operator.__lt__) @@ -162,21 +173,22 @@ def give_traits(*traits: Trait, exclude=None, override=None): to None. """ def wrapper(cls): - cls = dataclass(cls) nonlocal exclude nonlocal override if exclude is None: exclude = set() if override is None: override = dict() + annotations = dict(getattr(cls, "__annotations__", {})) for trait in traits: for attr, val in trait._fields.items(): if attr not in exclude: - if attr not in override: - setattr(cls, attr, val) - else: - setattr(cls, attr, override[attr]) - return cls + if attr in override: + val = override[attr] + annotations[attr] = type(val) + setattr(cls, attr, val) + cls.__annotations__ = annotations + return dataclass(cls) return wrapper @@ -245,8 +257,8 @@ def serialize(self, dest: str) -> None: """ serialized = {"entities" : []} for entity in self.entities: - entity_dict = entity.asdict() - entity_dict["obj_name"] = entity.__name__ + entity_dict = asdict(entity) + entity_dict["obj_name"] = type(entity).__name__ entity_dict["cls_name"] = type(entity).__name__ serialized["entities"].append(entity_dict) with open(dest, "w") as file: @@ -260,9 +272,24 @@ def deserialize(self, src: str): with open(src, "r") as file: entities = json.load(file) for entity in entities["entities"]: - initialized_entity = globals()[entity["cls_name"]] + initialized_entity = self._find_class(entity["cls_name"]) initialized_entity.__name__ = entity["obj_name"] - self.add(initialized_entity) + self.add(initialized_entity) + + def _find_class(self, cls_name: str) -> type: + """Look up a class by name in the caller's module, falling back to + all loaded modules.""" + frame = sys._getframe(2) + cls = frame.f_globals.get(cls_name) + if cls is None: + for module in sys.modules.values(): + candidate = getattr(module, cls_name, None) + if isinstance(candidate, type): + cls = candidate + break + if cls is None: + raise KeyError(cls_name) + return cls def extend(self, entities: list[Any]): """Add a collection of entities to the world. diff --git a/tests/test_basic.py b/tests/test_basic.py index a975677..9ff0617 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -1,6 +1,6 @@ # pyrefly: ignore-errors -from punyecs import World, Trait, requirements, one_shot, give_traits, ex_attr, has_attr, c, not_ +from punyecs import World, Trait, requirements, one_shot, give_traits, ex_attr, c, not_ def test_query(): diff --git a/tests/test_constraints.py b/tests/test_constraints.py new file mode 100644 index 0000000..746faaf --- /dev/null +++ b/tests/test_constraints.py @@ -0,0 +1,161 @@ +# pyrefly: ignore-errors + +from punyecs import ( + Query, + Trait, + and_, + c, + entity_satisfies_query, + ex_attr, + give_traits, + has_attr, + not_, + or_, +) + + +@give_traits(Trait(x=1.0, y=2.0, level=10, count=5, name="orc")) +class DummyEntity: + pass + + +def make_entity(): + return DummyEntity() + + +def test_attribute_access(): + e = make_entity() + c._obj = e + assert c.x.eval() == 1.0 + assert c.level.eval() == 10 + c._obj = None + + +def test_missing_attribute_evaluates_false(): + e = make_entity() + c._obj = e + assert c.missing.eval() is False + c._obj = None + + +def test_arithmetic_operators(): + e = make_entity() + c._obj = e + assert (c.x + 1).eval() == 2.0 + assert (c.x - 1).eval() == 0.0 + assert (c.x * 2).eval() == 2.0 + assert (c.x / 2).eval() == 0.5 + assert (c.count // 2).eval() == 2 + c._obj = None + + +def test_reflected_arithmetic_operators(): + e = make_entity() + c._obj = e + assert (1 + c.x).eval() == 2.0 + assert (2 * c.x).eval() == 2.0 + assert (1 - c.x).eval() == 0.0 + assert (2 / c.x).eval() == 2.0 + assert (10 // c.count).eval() == 2 + c._obj = None + + +def test_unary_operators(): + e = make_entity() + c._obj = e + assert (-c.x).eval() == -1.0 + assert (abs(c.x - 10)).eval() == 9.0 + assert (~c.count).eval() == -6 + c._obj = None + + +def test_comparison_operators(): + e = make_entity() + c._obj = e + assert (c.level > 5).eval() is True + assert (c.level < 5).eval() is False + assert (c.level == 10).eval() is True + assert (c.name == "orc").eval() is True + c._obj = None + + +def test_is_and_isnot(): + e = make_entity() + other = make_entity() + c._obj = e + assert c.is_(e).eval() is True + assert c.is_(other).eval() is False + assert c.isnot(e).eval() is False + assert c.isnot(other).eval() is True + c._obj = None + + +def test_not_(): + e = make_entity() + c._obj = e + assert not_(c.level > 5).eval() is False + assert not_(c.level > 100).eval() is True + c._obj = None + + +def test_bitwise_and_or_operators(): + e = make_entity() + c._obj = e + assert ((c.level > 5) & (c.count > 3)).eval() is True + assert ((c.level > 5) & (c.count > 100)).eval() is False + assert ((c.level > 100) | (c.count > 3)).eval() is True + assert ((c.level > 100) | (c.count > 100)).eval() is False + c._obj = None + + +def test_nested_arithmetic(): + e = make_entity() + c._obj = e + assert ((c.x + c.y) * 2).eval() == 6.0 + assert ((c.x + c.y) / 3).eval() == 1.0 + assert ((c.x + 2) * 2 > 5).eval() is True + c._obj = None + + +def test_has_attr_and_ex_attr(): + e = make_entity() + assert entity_satisfies_query(e, Query(Trait(), has_attr(c, "level"))) is True + assert entity_satisfies_query(e, Query(Trait(), has_attr(c, "missing"))) is False + assert entity_satisfies_query(e, Query(Trait(), ex_attr(c, "level"))) is False + assert entity_satisfies_query(e, Query(Trait(), ex_attr(c, "missing"))) is True + + +def test_entity_satisfies_query_attributes(): + e = make_entity() + assert entity_satisfies_query(e, Query(Trait(x=0.0, y=0.0))) is True + assert entity_satisfies_query(e, Query(Trait(x=0.0, y=0.0, missing=0.0))) is False + + +def test_entity_satisfies_query_constraint(): + e = make_entity() + assert entity_satisfies_query(e, Query(Trait(), c.level > 5)) is True + assert entity_satisfies_query(e, Query(Trait(), c.level > 100)) is False + + +def test_cursor_reset_after_query(): + e = make_entity() + entity_satisfies_query(e, Query(Trait(), c.level > 5)) + assert c._obj is None + entity_satisfies_query(e, Query(Trait(), c.level > 100)) + assert c._obj is None + + +def test_and_combines_constraints(): + e = make_entity() + c._obj = e + combined = and_(c.level > 5, c.count > 3) + assert combined.eval() is True + c._obj = None + + +def test_or_combines_constraints(): + e = make_entity() + c._obj = e + combined = or_(c.level > 100, c.count > 3) + assert combined.eval() is True + c._obj = None \ No newline at end of file diff --git a/tests/test_serialization.py b/tests/test_serialization.py new file mode 100644 index 0000000..a88aa94 --- /dev/null +++ b/tests/test_serialization.py @@ -0,0 +1,56 @@ +# pyrefly: ignore-errors + +import json + +from punyecs import Trait, World, give_traits, requirements + + +@give_traits(Trait(x=0.0, y=0.0)) +class SerializableEntity: + pass + + +def test_deserialize_loads_entities_from_json(tmp_path): + w = World() + + @requirements(w, Trait(x=0.0, y=0.0)) + def move(e, dt): + e.x += 0.1 + + data = { + "entities": [ + {"x": 0.0, "y": 0.0, "obj_name": "orc_1", "cls_name": "SerializableEntity"} + ] + } + src = tmp_path / "world.json" + src.write_text(json.dumps(data)) + + w.deserialize(str(src)) + + assert len(w.entities) == 1 + assert w.entities[0] is SerializableEntity + assert w.entities[0].__name__ == "orc_1" + assert len(w.groups) == 1 + assert w.groups[0].entities == [SerializableEntity] + + SerializableEntity.__name__ = "SerializableEntity" + + +def test_serialize_round_trip(tmp_path): + w = World() + entity = SerializableEntity() + w.add(entity) + + dest = tmp_path / "world.json" + w.serialize(str(dest)) + + assert json.loads(dest.read_text()) == { + "entities": [ + { + "x": 0.0, + "y": 0.0, + "obj_name": "SerializableEntity", + "cls_name": "SerializableEntity", + } + ] + } \ No newline at end of file diff --git a/tests/test_world.py b/tests/test_world.py new file mode 100644 index 0000000..279eeb2 --- /dev/null +++ b/tests/test_world.py @@ -0,0 +1,339 @@ +# pyrefly: ignore-errors + +import dataclasses + +from punyecs import Trait, World, c, give_traits, has_attr, one_shot, requirements + + +Pos = Trait(x=0.0, y=0.0) +Vel = Trait(vx=1.0, vy=2.0) +YAxis = Trait(y=0.0) +Stats = Trait(level=0) + + +def test_trait_combination(): + @give_traits(Pos + Vel) + class Mover: + pass + + m = Mover() + assert m.x == 0.0 + assert m.y == 0.0 + assert m.vx == 1.0 + assert m.vy == 2.0 + + +def test_multiple_traits(): + @give_traits(Pos, Vel) + class Mover: + pass + + m = Mover() + assert m.x == 0.0 + assert m.vx == 1.0 + + +def test_exclude_and_override_together(): + @give_traits(Pos, Vel, exclude={"y"}, override={"vx": 10.0}) + class Sprite: + pass + + s = Sprite() + assert s.x == 0.0 + assert not hasattr(s, "y") + assert s.vx == 10.0 + assert s.vy == 2.0 + + +def test_give_traits_makes_dataclass(): + @give_traits(Pos) + class Thing: + pass + + assert dataclasses.is_dataclass(Thing) + + +def test_update_respects_dt(): + w = World() + + @give_traits(Pos) + class Player: + pass + + @requirements(w, Pos) + def move(e, dt): + e.x += 0.5 * dt + + p = Player() + w.add(p) + w.update(2) + assert p.x == 1.0 + + +def test_systems_run_in_group_order(): + w = World() + + @give_traits(Pos) + class Enemy: + pass + + @requirements(w, Pos) + def add_one(e, dt): + e.x += 1 + + @requirements(w, Pos) + def double(e, dt): + e.x *= 2 + + e = Enemy() + w.add(e) + w.update(1) + assert e.x == 2.0 + + +def test_entity_added_after_system_declared(): + w = World() + + @give_traits(Pos) + class Enemy: + pass + + @requirements(w, Pos) + def move(e, dt): + e.x += 0.1 + + e = Enemy() + w.add(e) + w.update(1) + assert e.x == 0.1 + + +def test_add_entity_after_first_update(): + w = World() + + @give_traits(Pos) + class Enemy: + pass + + @requirements(w, Pos) + def move(e, dt): + e.x += 1 + + e1 = Enemy() + w.add(e1) + w.update(1) + e2 = Enemy() + w.add(e2) + w.update(1) + assert e1.x == 2 + assert e2.x == 1 + + +def test_extend(): + w = World() + + @give_traits(Pos) + class Enemy: + pass + + @requirements(w, Pos) + def move(e, dt): + e.x += 1 + + enemies = [Enemy() for _ in range(3)] + w.extend(enemies) + w.update(1) + assert [e.x for e in enemies] == [1, 1, 1] + + +def test_remove_multiple(): + w = World() + + @give_traits(Pos) + class Enemy: + pass + + @one_shot(w, Pos) + def inc_x(e): + e.x += 1 + + e1, e2, e3, e4 = Enemy(), Enemy(), Enemy(), Enemy() + w.extend([e1, e2, e3, e4]) + inc_x() + w.remove(e2, e4) + inc_x() + assert (e1.x, e2.x, e3.x, e4.x) == (2, 1, 2, 1) + + +def test_remove_unadded_entity_is_noop(): + w = World() + + @give_traits(Pos) + class Enemy: + pass + + @one_shot(w, Pos) + def inc_x(e): + e.x += 1 + + e1 = Enemy() + ghost = Enemy() + w.add(e1) + w.remove(ghost) + inc_x() + assert e1.x == 1 + + +def test_subject_to_level_filter(): + w = World() + + @give_traits(YAxis, Stats, override={"level": 30}) + class Grunt: + pass + + @give_traits(YAxis, Stats, override={"level": 80}) + class Elite: + pass + + @requirements(w, YAxis, subject_to=c.level > 50) + def gravity(e, dt): + e.y -= 1.0 * dt + + grunt = Grunt() + elite = Elite() + w.add(grunt) + w.add(elite) + w.update(1) + assert grunt.y == 0.0 + assert elite.y == -1.0 + + +def test_isnot_exclusion(): + w = World() + + @give_traits(Pos) + class Player: + pass + + @give_traits(Pos) + class Enemy: + pass + + player = Player() + enemy = Enemy() + + @requirements(w, Pos, subject_to=c.isnot(player)) + def move(e, dt): + e.x += 0.1 + e.y += 0.1 + + w.add(player) + w.add(enemy) + w.update(1) + assert player.x == 0.0 + assert player.y == 0.0 + assert enemy.x == 0.1 + assert enemy.y == 0.1 + + +def test_truthy_attribute_filter(): + w = World() + + @give_traits(Pos) + class Player: + controller: bool = True + + @give_traits(Pos) + class Enemy: + pass + + player = Player() + enemy = Enemy() + + @requirements(w, Pos, subject_to=c.controller) + def input(e, dt): + e.x += 1 + + w.add(player) + w.add(enemy) + w.update(1) + assert player.x == 1.0 + assert enemy.x == 0.0 + + +def test_has_attr_wiggle(): + w = World() + + @give_traits(Pos) + class Plain: + pass + + @give_traits(Pos) + class Wiggler: + wiggle = staticmethod(lambda x: x + 2) + + @requirements(w, Pos, subject_to=has_attr(c, "wiggle")) + def wiggle(e, dt): + e.x = e.wiggle(e.x) * dt + + plain = Plain() + wiggler = Wiggler() + w.add(plain) + w.add(wiggler) + w.update(1) + assert plain.x == 0.0 + assert wiggler.x == 2.0 + + +def test_one_shot_with_subject_to(): + w = World() + + @give_traits(Pos) + class Enemy: + pass + + e1 = Enemy() + e1.level = 10 + e2 = Enemy() + e2.level = 100 + + @one_shot(w, Pos, subject_to=c.level > 50) + def reward(e): + e.x += 10 + + w.add(e1) + w.add(e2) + reward() + assert e1.x == 0.0 + assert e2.x == 10.0 + + +def test_requirements_wraps_function(): + w = World() + + @give_traits(Pos) + class Enemy: + pass + + @requirements(w, Pos) + def move(e, dt): + e.x += dt + return e.x + + e = Enemy() + w.add(e) + result = move(e, 2) + assert result == 2.0 + + +def test_add_without_groups(): + w = World() + + @give_traits(Pos) + class Enemy: + pass + + e = Enemy() + w.add(e) + assert e in w.entities + assert w.groups == [] \ No newline at end of file