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