Fix children handling
This commit is contained in:
parent
603f8253cf
commit
75fd0e7bbc
1 changed files with 7 additions and 15 deletions
|
|
@ -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,9 +195,7 @@ class MemoryElementSupply(ElementSupply):
|
||||||
|
|
||||||
self._elements.pop(element_id)
|
self._elements.pop(element_id)
|
||||||
self._children.pop(element_id)
|
self._children.pop(element_id)
|
||||||
|
self._children[element.parent_id].remove(element.id)
|
||||||
if element.parent_id is not None:
|
|
||||||
self._children[element.parent_id].remove(element.id)
|
|
||||||
|
|
||||||
def get(self, element_id: Id) -> Element:
|
def get(self, element_id: Id) -> Element:
|
||||||
result = self._elements.get(element_id)
|
result = self._elements.get(element_id)
|
||||||
|
|
@ -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]]
|
result = self._children.get(element_id)
|
||||||
if element_id is None:
|
|
||||||
result = list(element.id for element in self._elements.values())
|
|
||||||
else:
|
|
||||||
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))
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue