Added and fixed tests.

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.

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():

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

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",
+            }
+        ]
+    }