gtn/ecs.py

68 lines
2.7 KiB
Python

from typing import Iterator
class World:
def __init__(self):
self._to_apply: set[Entity] = set()
self._entities: set[Entity] = set()
self._mapping: dict[type, set[Entity]] = {}
def create_entity(self, *components):
return Entity(self, *components)
def remove_entity(self, entity: 'Entity'):
entity._deleted = True
self._to_apply.add(entity)
def apply(self):
for entity in self._to_apply:
if entity._deleted:
self._entities.remove(entity)
for component_type in entity._components:
self._mapping[component_type].remove(entity)
else:
self._entities.add(entity)
for component_type, component in entity._to_apply:
if component is None:
del entity._components[component_type]
self._mapping[component_type].remove(entity)
else:
entity._components[component_type] = component
self._mapping.setdefault(component_type, set()).add(entity)
entity._to_apply.clear()
def query(self, *needed: type, without: tuple[type] = []) -> Iterator['Entity']:
if not needed:
for entity in self._entities:
if all(without_type not in entity for without_type in without):
yield entity
else:
for entity in self._mapping.get(needed[0], set()):
if all(entity in self._mapping.get(component_type, set()) for component_type in needed[1:]) and \
all(without_type not in entity for without_type in without):
yield entity
class Entity:
def __init__(self, world: World, *components):
self._world = world
self._to_apply: list[tuple[type, object]] = [(type(component), component) for component in components]
self._components: dict[type, object] = {}
self._deleted = False
self._world._to_apply.add(self)
def set(self, *components):
for component in components:
self._to_apply.append((type(component), component))
self._world._to_apply.add(self)
def remove(self, *component_types: type):
for component_type in component_types:
if component_type in self._components:
self._to_apply.append((component_type, None))
self._world._to_apply.add(self)
def __getitem__(self, component_type: type) -> object:
return self._components[component_type]
def __contains__(self, component_type: type) -> bool:
return component_type in self._components