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,
 34                                                  args={"component": component}))
 35
 36    def has_component(self, component: ComponentType) -> bool:
 37        """Checks if this entity has a component"""
 38        return component in self.get_components()
 39
 40    def get_components(self) -> list[ComponentType]:
 41        """get the components of this entity. Note: they may change, so call this every time, don't store it"""
 42        pool, _ = self._locate(names=[])
 43        return self._pool_to_components[pool]
 44
 45    def get_fields(self) -> set[str]:
 46        """gets the fields of this entity. Note: they may change, so call this every time, don't store it."""
 47        pool, _ = self._locate(names=[])
 48        return pool.fields_set
 49
 50    def set_data(self, **data):
 51        """Eagerly (no world.update()) sets the data of this entity. Multiple columns can be updated at once e.g.
 52        entity.set_data(a=x, b=y). Checks are done so object doesn't crash midway (e.g. bad dtype/shape etc.)."""
 53        pool, pool_index = self._locate(names=data.keys())
 54
 55        if len(data) == 1: # fast path because we don't need to do exhaustive tests on a single k->v data set
 56            k, v = next(iter(data.items()))
 57            pool.data[k][pool_index] = v
 58            return
 59
 60        # convert the data in numpy array of proper shape before calling pool.data[k][ix]=v so we don't have to revert
 61        ready: list[tuple[str, np.ndarray]] = []
 62        for k, v in data.items():
 63            col = pool.data[k]
 64            # no need to call the slow np.broadcast_to if v is already in proper shape and broadcastable and all.
 65            if not (isinstance(v, np.ndarray) and v.dtype == col.dtype and v.shape == col.shape[1:]):
 66                v = np.broadcast_to(np.asarray(v, dtype=col.dtype), col.shape[1:])
 67            ready.append((k, v))
 68
 69        # finally set the data after the conversion was done
 70        for k, v in ready:
 71            pool.data[k][pool_index] = v
 72
 73    def _locate(self, names: Iterable[str]) -> tuple[Pool, int]:
 74        try:
 75            pool, index = self._eid_to_pool_ix[self.entity_id]
 76        except KeyError:
 77            raise AttributeError(f"Entity {self.entity_id} not in world. Call `world.update()` if it was just added.")
 78
 79        if not (flds := pool.fields_set).issuperset(names):
 80            raise AttributeError(f"Not all of {list(names)} are fields (entity id: {self.entity_id}). "
 81                                 f"\n- Components: {[c.__name__ for c in self.get_components()]}\n- Fields: {flds}")
 82        return pool, index
 83
 84    def to_dict(self, serialization_field: str | None = None) -> dict[str, Any]:
 85        """
 86        Serializes a single entity. Assumes fields are numpy. numerics are converted via `.tolist()`. objects are
 87        converted via `.item()`.
 88        Parameters:
 89        - `serialization_field` An optional special field added at World-level (e.g.: 'serializable'). If set, then we
 90        only serialize this entity's fields where the serialization_field is True. If not set, all fields are dumped.
 91        """
 92        components = self.get_components()
 93        res = {"components": [c.__name__ for c in components], "data": {}}
 94        for component in components:
 95            for field in fields(component):
 96                # the magic key that we have added in extra_metadata at World level. If not set, all fields are dumped.
 97                if serialization_field is not None and field.metadata[serialization_field] is False:
 98                    continue
 99                if field.metadata["dtype"] == "object": # dtype=object is for... non-numeric data (mostly dicts)
100                    res["data"][field.name] = self.__getattr__(field.name).item()
101                else:
102                    res["data"][field.name] = self.__getattr__(field.name).tolist()
103        return res
104
105    def __getattr__(self, name: str) -> np.ndarray:
106        try: # note: try/catch is cheaper than self._locate(name=[names])
107            pool, index = self._eid_to_pool_ix[self.entity_id]
108            return pool.data[name][index]
109        except KeyError:
110            if self.entity_id not in self._eid_to_pool_ix:
111                raise AttributeError(f"Entity {self.entity_id} not in world. Perhaps call `world.update()`.")
112            raise AttributeError(f"Entity {self.entity_id} may not have field: {name} (fields: {self.get_fields()})")
113
114    def __setattr__(self, name: str, value: np.ndarray):
115        if name in _ENTITY_INTERNAL_ATTRS:
116            super().__setattr__(name, value)
117            return
118
119        try: # note: try/catch is cheaper than self._locate(name=[names])
120            pool, index = self._eid_to_pool_ix[self.entity_id]
121            pool.data[name][index] = value
122        except KeyError:
123            if self.entity_id not in self._eid_to_pool_ix:
124                raise AttributeError(f"Entity {self.entity_id} not in world. Perhaps call `world.update()`.")
125            raise AttributeError(f"Entity {self.entity_id} may not have field: {name} (fields: {self.get_fields()})")
126
127    def __reduce__(self):
128        raise TypeError("Entity is a live view into the world's pools; it cannot be copied or pickled. "
129                        "Use entity.to_dict(), or serialize the world.")
130
131    def __repr__(self):
132        return f"EID-{self.entity_id}"
133
134ENTITY_RESERVED_NAMES = _ENTITY_INTERNAL_ATTRS | {n for n in vars(Entity) if not n.startswith("__")}