fix remove_entity method

1 files changed, 7 insertions, 5 deletions

diff --git a/kennel/engine/entities/entity.py b/kennel/engine/entities/entity.py
index b72e996..6623512 100644
--- a/kennel/engine/entities/entity.py
+++ b/kennel/engine/entities/entity.py
@@ -57,11 +57,13 @@ class EntityManager:
             self.component_entities[component.component_type].add(entity)
 
     def remove_entity(self, entity_id: str) -> None:
-        if entity_id in self.entities:
-            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)
+        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)