Fix children handling

This commit is contained in:
Joscha 2019-05-12 08:50:32 +00:00
parent 603f8253cf
commit 75fd0e7bbc

View file

@ -1,5 +1,5 @@
import abc import abc
from typing import Dict, List, Optional from typing import Dict, List, Optional, Set
from .element import Element, Id from .element import Element, Id
from .exceptions import TreeException from .exceptions import TreeException
@ -170,7 +170,7 @@ class MemoryElementSupply(ElementSupply):
def __init__(self) -> None: def __init__(self) -> None:
self._elements: Dict[Id, Element] = {} self._elements: Dict[Id, Element] = {}
self._children: Dict[Id, List[Id]] = {} self._children: Dict[Optional[Id], Set[Id]] = {None: set()}
def add(self, element: Element) -> None: def add(self, element: Element) -> None:
""" """
@ -181,10 +181,8 @@ class MemoryElementSupply(ElementSupply):
self.remove(element.id) self.remove(element.id)
self._elements[element.id] = element self._elements[element.id] = element
self._children[element.id] = [] self._children[element.id] = set()
self._children[element.parent_id].add(element.id)
if element.parent_id is not None:
self._children[element.parent_id].append(element.id)
def remove(self, element_id: Id) -> None: def remove(self, element_id: Id) -> None:
""" """
@ -197,8 +195,6 @@ class MemoryElementSupply(ElementSupply):
self._elements.pop(element_id) self._elements.pop(element_id)
self._children.pop(element_id) self._children.pop(element_id)
if element.parent_id is not None:
self._children[element.parent_id].remove(element.id) self._children[element.parent_id].remove(element.id)
def get(self, element_id: Id) -> Element: def get(self, element_id: Id) -> Element:
@ -210,13 +206,9 @@ class MemoryElementSupply(ElementSupply):
return result return result
def get_children_ids(self, element_id: Optional[Id]) -> List[Id]: def get_children_ids(self, element_id: Optional[Id]) -> List[Id]:
result: Optional[List[Id]]
if element_id is None:
result = list(element.id for element in self._elements.values())
else:
result = self._children.get(element_id) result = self._children.get(element_id)
if result is None: if result is None:
raise TreeException(f"Element with id {element_id!r} could not be found") raise TreeException(f"Element with id {element_id!r} could not be found")
return result return list(sorted(result))