microecs.entity

  1"""entity.py - A view of an entity with all its fields from the pool it belongs to in the world"""
  2from __future__ import annotations
  3from dataclasses import fields
  4from typing import Any, Iterable
  5import numpy as np
  6from .pool import Pool
  7from .component import ComponentType
  8from .utils import EntityId
  9from .command_buffer import CommandBuffer, Command, CommandType
 10
 11# Note: if Entity gets new fields, add them here! Otherwise the user code may overwrite them e.g. ent._eid_to_pool_ix=xx
 12_ENTITY_INTERNAL_ATTRS = {"entity_id", "_eid_to_pool_ix", "_pool_to_components", "_world_command_buffer"}
 13
 14class Entity:
 15    """
 16    A view of an entity with all its fields from the pool it belongs to in the world.
 17    Note: Consistent to internal pool changes, however it always must check where it belongs so it's slow!!
 18    """
 19    def __init__(self, entity_id: EntityId, eid_to_pool_ix: dict[EntityId, tuple[Pool, int]],
 20                 pool_to_components: dict[Pool, list[ComponentType]], world_command_buffer: CommandBuffer):
 21        self.entity_id = entity_id
 22        self._eid_to_pool_ix = eid_to_pool_ix
 23        self._pool_to_components = pool_to_components
 24        self._world_command_buffer = world_command_buffer # the world command buffer, needed for add/remove_component
 25
 26    def add_component(self, component: ComponentType, **kwargs):
 27        """Adds a component to an entity. Component data is sent to kwargs. Lazy; call world.update()"""
 28        self._world_command_buffer.append(Command(CommandType.ADD_COMPONENT, self.entity_id,
 29                                          args={"component": component, **kwargs}))
 30
 31    def remove_component(self, component: ComponentType):
 32        """Removes a component from an entity given its id. Lazy; call update()"""
 33        self._world_command_buffer.append(Command(CommandType.REMOVE_COMPONENT, self.entity_id, args=component))
 34
 35    def has_component(self, component: ComponentType) -> bool:
 36        """Checks if this entity has a component"""
 37        return component in self.get_components()
 38
 39    def get_components(self) -> list[ComponentType]:
 40        """get the components of this entity. Note: they may change, so call this every time, don't store it"""
 41        pool, _ = self._locate(names=[])
 42        return self._pool_to_components[pool]
 43
 44    def get_fields(self) -> set[str]:
 45        """gets the fields of this entity. Note: they may change, so call this every time, don't store it."""
 46        pool, _ = self._locate(names=[])
 47        return pool.fields_set
 48
 49    def set_data(self, **data):
 50        """Eagerly (no world.update()) sets the data of this entity. Multiple columns can be updated at once e.g.
 51        entity.set_data(a=x, b=y). Checks are done so object doesn't crash midway (e.g. bad dtype/shape etc.)."""
 52        pool, pool_index = self._locate(names=data.keys())
 53
 54        if len(data) == 1: # fast path because we don't need to do exhaustive tests on a single k->v data set
 55            k, v = next(iter(data.items()))
 56            pool.data[k][pool_index] = v
 57            return
 58
 59        # convert the data in numpy array of proper shape before calling pool.data[k][ix]=v so we don't have to revert
 60        ready: list[tuple[str, np.ndarray]] = []
 61        for k, v in data.items():
 62            col = pool.data[k]
 63            # no need to call the slow np.broadcast_to if v is already in proper shape and broadcastable and all.
 64            if not (isinstance(v, np.ndarray) and v.dtype == col.dtype and v.shape == col.shape[1:]):
 65                v = np.broadcast_to(np.asarray(v, dtype=col.dtype), col.shape[1:])
 66            ready.append((k, v))
 67
 68        # finally set the data after the conversion was done
 69        for k, v in ready:
 70            pool.data[k][pool_index] = v
 71
 72    def _locate(self, names: Iterable[str]) -> tuple[Pool, int]:
 73        try:
 74            pool, index = self._eid_to_pool_ix[self.entity_id]
 75        except KeyError:
 76            raise AttributeError(f"Entity {self.entity_id} not in world. Call `world.update()` if it was just added.")
 77
 78        if not (flds := pool.fields_set).issuperset(names):
 79            raise AttributeError(f"Not all of {list(names)} are fields (entity id: {self.entity_id}). "
 80                                 f"\n- Components: {[c.__name__ for c in self.get_components()]}\n- Fields: {flds}")
 81        return pool, index
 82
 83    def to_dict(self, serialization_field: str | None = None) -> dict[str, Any]:
 84        """
 85        Serializes a single entity. Assumes fields are numpy. numerics are converted via `.tolist()`. objects are
 86        converted via `.item()`.
 87        Parameters:
 88        - `serialization_field` An optional special field added at World-level (e.g.: 'serializable'). If set, then we
 89        only serialize this entity's fields where the serialization_field is True. If not set, all fields are dumped.
 90        """
 91        components = self.get_components()
 92        res = {"components": [c.__name__ for c in components], "data": {}}
 93        for component in components:
 94            for field in fields(component):
 95                # the magic key that we have added in extra_metadata at World level. If not set, all fields are dumped.
 96                if serialization_field is not None and field.metadata[serialization_field] is False:
 97                    continue
 98                if field.metadata["dtype"] == "object": # dtype=object is for... non-numeric data (mostly dicts)
 99                    res["data"][field.name] = self.__getattr__(field.name).item()
100                else:
101                    res["data"][field.name] = self.__getattr__(field.name).tolist()
102        return res
103
104    def __getattr__(self, name: str) -> np.ndarray:
105        try: # note: try/catch is cheaper than self._locate(name=[names])
106            pool, index = self._eid_to_pool_ix[self.entity_id]
107            return pool.data[name][index]
108        except KeyError:
109            if self.entity_id not in self._eid_to_pool_ix:
110                raise AttributeError(f"Entity {self.entity_id} not in world. Perhaps call `world.update()`.")
111            raise AttributeError(f"Entity {self.entity_id} may not have field: {name} (fields: {self.get_fields()})")
112
113    def __setattr__(self, name: str, value: np.ndarray):
114        if name in _ENTITY_INTERNAL_ATTRS:
115            super().__setattr__(name, value)
116            return
117
118        try: # note: try/catch is cheaper than self._locate(name=[names])
119            pool, index = self._eid_to_pool_ix[self.entity_id]
120            pool.data[name][index] = value
121        except KeyError:
122            if self.entity_id not in self._eid_to_pool_ix:
123                raise AttributeError(f"Entity {self.entity_id} not in world. Perhaps call `world.update()`.")
124            raise AttributeError(f"Entity {self.entity_id} may not have field: {name} (fields: {self.get_fields()})")
125
126    def __reduce__(self):
127        raise TypeError("Entity is a live view into the world's pools; it cannot be copied or pickled. "
128                        "Use entity.to_dict(), or serialize the world.")
129
130    def __repr__(self):
131        return f"EID-{self.entity_id}"
132
133ENTITY_RESERVED_NAMES = _ENTITY_INTERNAL_ATTRS | {n for n in vars(Entity) if not n.startswith("__")}