222 lines
6.5 KiB
Python
222 lines
6.5 KiB
Python
import abc
|
|
from typing import Dict, List, Optional
|
|
|
|
from .element import Element, Id
|
|
from .exceptions import TreeException
|
|
|
|
__all__ = ["ElementSupply", "MemoryElementSupply"]
|
|
|
|
class ElementSupply(abc.ABC):
|
|
"""
|
|
An ElementSupply is an interface to query some resource containing
|
|
Elements. The elements could for example be kept in memory, in a database
|
|
or somewhere else.
|
|
|
|
The element ids must be unique, and the elements and their parents must
|
|
form one or more trees (i. e. must not contain any cycles).
|
|
"""
|
|
|
|
@abc.abstractmethod
|
|
def get(self, element_id: Id) -> Element:
|
|
"""
|
|
Get a single element by its id.
|
|
"""
|
|
|
|
pass
|
|
|
|
def get_parent_id(self, element_id: Id) -> Optional[Id]:
|
|
"""
|
|
Get the id of the parent's element.
|
|
|
|
This function is redundant, since you can just use element.parent_id.
|
|
"""
|
|
|
|
return self.get(element_id).parent_id
|
|
|
|
def get_parent(self, element_id: Id) -> Optional[Element]:
|
|
"""
|
|
Like get_parent_id, but returns the Element instead.
|
|
"""
|
|
|
|
parent_id = self.get_parent_id(element_id)
|
|
|
|
if parent_id is not None:
|
|
return self.get(parent_id)
|
|
else:
|
|
return None
|
|
|
|
@abc.abstractmethod
|
|
def get_children_ids(self, element_id: Optional[Id]) -> List[Id]:
|
|
"""
|
|
Get a list of the ids of all the element's children.
|
|
"""
|
|
|
|
pass
|
|
|
|
def get_children(self, element_id: Optional[Id]) -> List[Element]:
|
|
"""
|
|
Get a list of all children of an element.
|
|
|
|
If the id passed is None, return a list of all top-level elements
|
|
instead.
|
|
"""
|
|
|
|
children_ids = self.get_children_ids(element_id)
|
|
|
|
children: List[Element] = []
|
|
for child_id in children_ids:
|
|
children.append(self.get(child_id))
|
|
|
|
return children
|
|
|
|
# There is not a clear-cut way to get the previous or next "sibling" of an
|
|
# element that is itself a child of the implicit root (None), since
|
|
# get_children() doesn't accept None in its element_id argument.
|
|
#
|
|
# Because of this, the get_previous_id() and get_next_id() functions are
|
|
# abstract (until I decide to change the signature of get_children(), that
|
|
# is :P).
|
|
|
|
def get_previous_id(self, element_id: Id) -> Optional[Id]:
|
|
"""
|
|
Get the id of an element's previous sibling (i. e. the sibling just
|
|
above it).
|
|
|
|
Returns None if there is no previous sibling.
|
|
|
|
Depending on the amount of elements in your ElementSupply, the default
|
|
implementation might get very slow and/or use a lot of memory.
|
|
"""
|
|
|
|
siblings = self.get_children_ids(self.get_parent_id(element_id))
|
|
index = siblings.index(element_id)
|
|
|
|
if index <= 0:
|
|
return None
|
|
else:
|
|
return siblings[index - 1]
|
|
|
|
def get_previous(self, element_id: Id) -> Optional[Element]:
|
|
"""
|
|
Like get_previous_id(), but returns the Element instead.
|
|
"""
|
|
|
|
previous_id = self.get_previous_id(element_id)
|
|
|
|
if previous_id is not None:
|
|
return self.get(previous_id)
|
|
else:
|
|
return None
|
|
|
|
def get_next_id(self, element_id: Id) -> Optional[Id]:
|
|
"""
|
|
Get the id of an element's next sibling (i. e. the sibling just below
|
|
it).
|
|
|
|
Returns None if there is no next sibling.
|
|
|
|
Depending on the amount of elements in your ElementSupply, the default
|
|
implementation might get very slow and/or use a lot of memory.
|
|
"""
|
|
|
|
siblings = self.get_children_ids(self.get_parent_id(element_id))
|
|
index = siblings.index(element_id)
|
|
|
|
if index >= len(siblings) - 1:
|
|
return None
|
|
else:
|
|
return siblings[index + 1]
|
|
|
|
def get_next(self, element_id: Id) -> Optional[Element]:
|
|
"""
|
|
Like get_next_id(), but returns the Element instead.
|
|
"""
|
|
|
|
next_id = self.get_next_id(element_id)
|
|
|
|
if next_id is not None:
|
|
return self.get(next_id)
|
|
else:
|
|
return None
|
|
|
|
# def get_furthest_ancestor_id(self,
|
|
# element_id: Id,
|
|
# root_id: Optional[Id] = None,
|
|
# ) -> Id:
|
|
# current_id = element_id
|
|
#
|
|
# while True:
|
|
# parent_id = self.get_parent_id(current_id)
|
|
#
|
|
# if parent_id == root_id:
|
|
# return current_id
|
|
# elif parent_id is None:
|
|
# raise TreeException("Reached implicit root before hitting specified root")
|
|
#
|
|
# current_id = parent_id
|
|
#
|
|
# def get_furthest_ancestor(self,
|
|
# element_id: Id,
|
|
# root_id: Optional[Id] = None,
|
|
# ) -> Element:
|
|
# return self.get(self.get_furthest_ancestor_id(element_id,
|
|
# root_id=root_id))
|
|
|
|
class MemoryElementSupply(ElementSupply):
|
|
"""
|
|
An in-memory implementation of an ElementSupply that works with any type of
|
|
Element.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._elements: Dict[Id, Element] = {}
|
|
self._children: Dict[Id, List[Id]] = {}
|
|
|
|
def add(self, element: Element) -> None:
|
|
"""
|
|
Add a new element or overwrite an existing element with the same id.
|
|
"""
|
|
|
|
if element.id in self._elements:
|
|
self.remove(element.id)
|
|
|
|
self._elements[element.id] = element
|
|
self._children[element.id] = []
|
|
|
|
if element.parent_id is not None:
|
|
self._children[element.parent_id].append(element.id)
|
|
|
|
def remove(self, element_id: Id) -> None:
|
|
"""
|
|
Remove an element. This function does nothing if the element doesn't
|
|
exist in this ElementSupply.
|
|
"""
|
|
|
|
if element_id in self._elements:
|
|
element = self.get(element_id)
|
|
|
|
self._elements.pop(element_id)
|
|
self._children.pop(element_id)
|
|
|
|
if element.parent_id is not None:
|
|
self._children[element.parent_id].remove(element.id)
|
|
|
|
def get(self, element_id: Id) -> Element:
|
|
result = self._elements.get(element_id)
|
|
|
|
if result is None:
|
|
raise TreeException(f"Element with id {element_id!r} could not be found")
|
|
|
|
return result
|
|
|
|
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)
|
|
|
|
if result is None:
|
|
raise TreeException(f"Element with id {element_id!r} could not be found")
|
|
|
|
return result
|