73 lines
2.5 KiB
Python
73 lines
2.5 KiB
Python
from enum import Enum
|
|
from kennel.engine.components.component import Component, ComponentType
|
|
from typing import List, Optional
|
|
|
|
|
|
class EntityType(str, Enum):
|
|
LASER = "LASER"
|
|
|
|
|
|
class Entity:
|
|
def __init__(
|
|
self, entity_type: EntityType, id: str, component_list: List[Component]
|
|
):
|
|
self.entity_type = entity_type
|
|
self.id = id
|
|
self.components = {}
|
|
for component in component_list:
|
|
self.add_component(component)
|
|
|
|
def get_component(self, component_type: ComponentType) -> Optional[Component]:
|
|
return self.components[component_type]
|
|
|
|
def add_component(self, component: Component) -> None:
|
|
self.components[component.component_type] = component
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"entity_type": self.entity_type,
|
|
"id": self.id,
|
|
"components": {k: v.dict() for k, v in self.components.items()},
|
|
}
|
|
|
|
|
|
class EntityManager:
|
|
def __init__(self):
|
|
self.entities = {}
|
|
self.component_entities = {}
|
|
|
|
def update(self) -> None:
|
|
self.component_entities = {}
|
|
for entity in self.entities.values():
|
|
for component in entity.components.values():
|
|
if component.component_type not in self.component_entities:
|
|
self.component_entities[component.component_type] = set()
|
|
self.component_entities[component.component_type].add(entity)
|
|
|
|
def get_entities_with_component(
|
|
self, component_type: ComponentType
|
|
) -> List[Entity]:
|
|
return self.component_entities.get(component_type, [])
|
|
|
|
def add_entity(self, entity: Entity) -> None:
|
|
self.entities[entity.id] = entity
|
|
for component in entity.components.values():
|
|
if component.component_type not in self.component_entities:
|
|
self.component_entities[component.component_type] = set()
|
|
self.component_entities[component.component_type].add(entity)
|
|
|
|
def remove_entity(self, entity_id: str) -> None:
|
|
if entity_id not in self.entities:
|
|
return
|
|
entity = self.entities[entity_id]
|
|
for component in entity.components.values():
|
|
if component.component_type in self.component_entities:
|
|
self.component_entities[component.component_type].remove(entity)
|
|
del self.entities[entity_id]
|
|
|
|
def get_entity(self, entity_id: str) -> Optional[Entity]:
|
|
return self.entities.get(entity_id)
|
|
|
|
def to_dict(self) -> dict:
|
|
return {k: v.to_dict() for k, v in self.entities.items()}
|