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]] = {}
        self._to_apply_resources: list[tuple[type, object]] = []
        self._resources: dict[type, object] = {}
    
    def create_entity(self, *components):
        return Entity(self, *components)

    def remove_entity(self, entity: 'Entity'):
        entity._deleted = True
        self._to_apply.add(entity)
    
    def set(self, *resources):
        for resource in resources:
            self._to_apply_resources.append((type(resource), resource))
    
    def remove(self, *resource_types):
        for resource_type in resource_types:
            if resource_type in self._resources:
                self._to_apply_resources.append((resource_type, None))

    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()
        self._to_apply.clear()
        for resource_type, resource in self._to_apply_resources:
            if resource is None: del self._resources[resource_type]
            else: self._resources[resource_type] = resource
        self._to_apply_resources.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
    
    def __getitem__(self, resource_type: type) -> object:
        return self._resources[resource_type]
    
    def __contains__(self, resource_type: type) -> bool:
        return resource_type in self._resources

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