gtn/ecs.py

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