Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 22 additions & 29 deletions lighttree/tree.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,11 @@
from .utils import STYLES


Key = Union[None, str, int]
KeyedNode = Tuple[Key, Node]
KeyedTree = Tuple[Key, "Tree"]


class Tree(object):

"""Principles:
Expand Down Expand Up @@ -57,9 +62,7 @@ def __init__(self, path_separator: str = ".") -> None:
def __contains__(self, identifier: str) -> bool:
return identifier in self._nodes_map

def get(
self, nid: str, by_path: bool = False
) -> Tuple[Union[None, str, int], Node]:
def get(self, nid: str, by_path: bool = False) -> KeyedNode:
"""Get a node by its id.
:param nid: str, identifier of node to fetch
:param by_path: bool, if True nid is the path to the node
Expand All @@ -86,9 +89,7 @@ def child_id(self, nid: str, key: Union[str, int], by_path: bool = False) -> str
raise ValueError("Expected integer key, got %s" % key)
return self._nodes_children_list[nid][int(key)]

def child(
self, nid: str, key: Union[int, str], by_path: bool = False
) -> Tuple[Union[None, str, int], Node]:
def child(self, nid: str, key: Union[int, str], by_path: bool = False) -> KeyedNode:
return self.get(self.child_id(nid, key, by_path=by_path))

def get_node_id_by_path(self, path: str) -> str:
Expand All @@ -114,7 +115,7 @@ def get_path(self, nid: str) -> str:
]
)

def get_key(self, nid: str) -> Union[int, str, None]:
def get_key(self, nid: str) -> Key:
"""Get a node's key.
:param nid: str, identifier of node

Expand All @@ -135,7 +136,7 @@ def list(
id_in: Optional[Sequence[str]] = None,
depth_in: Optional[Sequence[int]] = None,
filter_: Optional[Callable[[Node], bool]] = None,
) -> List[Tuple[Union[None, int, str], Node]]:
) -> List[KeyedNode]:
"""List nodes.
:param id_in: list of str, optional, filter nodes among provided identifiers
:param depth_in: list of int, optional, filter nodes whose depth in tree is among provided values
Expand Down Expand Up @@ -234,7 +235,7 @@ def clone(
new_tree.insert_node(node, parent_id=pid, key=key)
return new_tree

def parent(self, nid: str) -> Tuple[Union[None, str, int], Node]:
def parent(self, nid: str) -> KeyedNode:
"""Return parent node.
Return None if given node id is root.
"""
Expand All @@ -255,9 +256,7 @@ def parent_id(self, nid: str, by_path: bool = False) -> str:
raise NotFoundNodeError()
return parent_id

def children(
self, nid: str, by_path: bool = False
) -> List[Tuple[Union[None, str, int], Node]]:
def children(self, nid: str, by_path: bool = False) -> List[KeyedNode]:
"""Return set of given node children node ids."""
return [self.get(id_) for id_ in self.children_ids(nid, by_path=by_path)]

Expand All @@ -266,9 +265,7 @@ def children_ids(self, nid: str, by_path: bool = False) -> List[str]:
return list(self._nodes_children_map[nid].keys())
return list(self._nodes_children_list[nid])

def siblings(
self, nid: str, by_path: bool = False
) -> List[Tuple[Union[None, str, int], Node]]:
def siblings(self, nid: str, by_path: bool = False) -> List[KeyedNode]:
"""Return set of ids of nodes that share the provided node's parent."""
return [self.get(id_) for id_ in self.siblings_ids(nid, by_path=by_path)]

Expand Down Expand Up @@ -297,7 +294,7 @@ def ancestors(
from_root: bool = False,
include_current: bool = False,
by_path: bool = False,
) -> List[Tuple[Union[None, str, int], Node]]:
) -> List[KeyedNode]:
"""From element to root.
:param nid:
:param from_root:
Expand Down Expand Up @@ -332,9 +329,7 @@ def ancestors_ids(
ancestor_ids = list(reversed(ancestor_ids))
return ancestor_ids

def subtree(
self, nid: str, deep: bool = False, by_path: bool = False
) -> Tuple[Union[None, str, int], "Tree"]:
def subtree(self, nid: str, deep: bool = False, by_path: bool = False) -> KeyedTree:
if by_path:
nid = self.get_node_id_by_path(nid)
t = self.clone(with_nodes=True, new_root=nid, deep=deep)
Expand All @@ -344,7 +339,7 @@ def subtree(

def leaves(
self, nid: Optional[str] = None, by_path: bool = False
) -> List[Tuple[Union[None, str, int], Node]]:
) -> List[KeyedNode]:
"""Return leaves under a node subtree."""
return [self.get(id_) for id_ in self.leaves_ids(nid, by_path=by_path)]

Expand Down Expand Up @@ -395,7 +390,7 @@ def insert_node(
child_id: Optional[str] = None,
key: Optional[Union[int, str]] = None,
by_path: bool = False,
) -> Union[None, int, str]:
) -> Key:
"""Insert node, return key
:param node:
:param parent_id:
Expand Down Expand Up @@ -498,7 +493,7 @@ def insert_tree(
child_id_below: Optional[str] = None,
key: Optional[Union[str, int]] = None,
by_path: bool = False,
) -> Union[int, str, None]:
) -> Key:
"""Return new key"""
self._validate_tree_insertion(new_tree)
if new_tree.root is None:
Expand Down Expand Up @@ -580,7 +575,7 @@ def _insert_tree_above(
self._insert_tree_below(new_tree, parent_id, key=subtree_key, by_path=False)
self._insert_tree_below(child_subtree, child_id_below, key=key, by_path=False)

def _drop_node(self, nid: str) -> Tuple[Union[None, str, int], Node]:
def _drop_node(self, nid: str) -> KeyedNode:
"""Return key, node"""
if self.children_ids(nid):
raise ValueError("Cannot drop node having children.")
Expand Down Expand Up @@ -611,7 +606,7 @@ def drop_node(
nid: str,
with_children: bool = True,
by_path: bool = False,
) -> Tuple[Union[None, str, int], Node]:
) -> KeyedNode:
"""If with_children is False, children of this node will take as new parent the dropped node parent.
Possible only if node type is same as parent node type.

Expand Down Expand Up @@ -645,9 +640,7 @@ def drop_node(
self._insert_tree_below(new_tree=st, parent_id=pid, key=k, by_path=False)
return removed_key, node

def drop_subtree(
self, nid: str, by_path: bool = False
) -> Tuple[Union[None, str, int], "Tree"]:
def drop_subtree(self, nid: str, by_path: bool = False) -> KeyedTree:
if by_path:
nid = self.get_node_id_by_path(nid)
self._ensure_present(nid)
Expand All @@ -663,7 +656,7 @@ def expand_tree(
filter_: Optional[Callable[[Union[None, str, int], Node], bool]] = None,
filter_through: bool = False,
reverse: bool = False,
) -> Iterable[Tuple[Union[None, int, str], Node]]:
) -> Iterable[KeyedNode]:
"""Python generator traversing the tree (or a subtree) with optional node filtering.

Inspired by treelib implementation https://github.com/caesar0301/treelib/blob/master/treelib/tree.py#L374
Expand Down Expand Up @@ -781,7 +774,7 @@ def _iter_nodes_with_location(
filter_: Optional[Callable[[Node], bool]],
reverse: bool,
is_last_list: Optional[List[bool]] = None,
) -> Iterable[Tuple[Tuple[bool, ...], Union[None, int, str], Node]]:
) -> Iterable[Tuple[Tuple[bool, ...], Key, Node]]:
"""Yield nodes with information on how they are placed.
:param nid: starting node identifier
:param filter_: filter function applied on nodes
Expand Down
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@

from setuptools import setup

__version__ = "1.1.0"
__version__ = "1.1.1"

develop_requires = [
"pre-commit",
Expand Down