Initial commit

ecs.py

@@ -0,0 +1,57 @@
+from dataclasses import dataclass, field
+from typing import Any, Iterator, Protocol
+
+type Entity = int
+type Archetype = frozenset[type]
+type ComponentMap = dict[type, Any]
+type EntityComponentStorage = dict[Entity, ComponentMap]
+
+
+class System(Protocol):
+    @classmethod
+    def process(cls, world: "World") -> None: ...
+
+
+@dataclass
+class World:
+    registry: dict[Archetype, EntityComponentStorage] = field(default_factory=dict)
+    next_entity_index: int = 0
+    systems: list[type[System]] = field(default_factory=list)
+
+    def generate_id(self) -> int:
+        self.next_entity_index += 1
+        return self.next_entity_index
+
+    def spawn(self, *components) -> int:
+        archetype = frozenset([type(c) for c in components])
+
+        if archetype not in self.registry:
+            self.registry[archetype] = {}
+
+        entity = self.generate_id()
+        self.registry[archetype][entity] = {type(c): c for c in components}
+        return entity
+
+    def query(
+        self, *component_types, exact=False
+    ) -> Iterator[tuple[Entity, ComponentMap]]:
+        component_subset_signature = frozenset(component_types)
+        for archetype, entity_component_storage in self.registry.items():
+            if (exact and component_subset_signature == archetype) or (
+                not exact and component_subset_signature.issubset(archetype)
+            ):
+                for entity, component_storage in entity_component_storage.items():
+                    yield entity, component_storage
+
+    def get_components(self, entity: int, pop=False):
+        for componentstorage in self.registry.values():
+            if entity in componentstorage:
+                if pop:
+                    return componentstorage.pop(entity, None)
+                else:
+                    return componentstorage[entity]
+        return None
+
+    def tick(self):
+        for system in self.systems:
+            system.process(self)