usage of generic Node in Tree definition, dataclass Node

Author: leonardbinet

lighttree/__init__.py

@@ -1,4 +1,13 @@
 from .interactive import TreeBasedObj
-from .tree import Tree, Node, Key, KeyedNode, KeyedTree
+from .tree import Tree, Key, KeyedNode, KeyedTree
+from .node import Node, AutoIdNode
 
-__all__ = ["Tree", "Node", "TreeBasedObj", "Key", "KeyedNode", "KeyedTree"]
+__all__ = [
+    "Tree",
+    "Node",
+    "AutoIdNode",
+    "TreeBasedObj",
+    "Key",
+    "KeyedNode",
+    "KeyedTree",
+]

lighttree/implementations/json_tree.py

@@ -1,5 +1,5 @@
 from typing import Dict, Optional, Any, Union, List
-from lighttree import Node, Tree, Key
+from lighttree import Tree, Key, AutoIdNode
 from lighttree.node import NodeId
 from lighttree.interactive import TreeBasedObj
 
@@ -32,26 +32,26 @@ def _fill(self, data: Any, key: Key, strict: bool, path: str = "") -> None:
         else:
             pid = self.get_node_id_by_path(path=path)
         if isinstance(data, list) or not strict and isinstance(data, tuple):
-            k = self.insert_node(Node(keyed=False), parent_id=pid, key=key)
+            k = self.insert_node(AutoIdNode(keyed=False), parent_id=pid, key=key)
             path = self._concat(path, k)
             for el in data:
                 self._fill(el, strict=strict, path=path, key=None)
             return
         if isinstance(data, dict):
-            k = self.insert_node(Node(keyed=True), key=key, parent_id=pid)
+            k = self.insert_node(AutoIdNode(keyed=True), key=key, parent_id=pid)
             path = self._concat(path, k)
             for sk, el in data.items():
                 self._fill(el, strict=strict, path=path, key=sk)
             return
         if isinstance(data, (str, int, float)):
             self.insert_node(
-                Node(accept_children=False, repr_=str(data), data=data),
+                AutoIdNode(accept_children=False, repr_=str(data), data=data),
                 parent_id=pid,
                 key=key,
             )
             return
         if data is None:
-            self.insert_node(Node(accept_children=False), parent_id=pid)
+            self.insert_node(AutoIdNode(accept_children=False), parent_id=pid)
             return
         raise TypeError("Unsupported type %s" % type(data))
 

lighttree/node.py

@@ -1,54 +1,63 @@
 import uuid
 from typing import Optional, Any, Tuple
+from dataclasses import dataclass
 
 NodeId = str
 
 
-class Node(object):
-    def __init__(
-        self,
-        identifier: Optional[NodeId] = None,
-        auto_uuid: bool = True,
-        keyed: bool = True,
-        accept_children: bool = True,
-        repr_: Optional[str] = None,
-        data: Any = None,
-    ) -> None:
-        """
-        :param identifier: node identifier, must be unique per tree
-        """
-        if identifier is None:
-            if not auto_uuid:
-                raise ValueError("Required identifier")
-            identifier = str(uuid.uuid4())
-        self.identifier = identifier
-        self.keyed = keyed
-        self.accept_children = accept_children
-        self.repr = repr_
-        self.data = data
+@dataclass
+class Node:
+
+    identifier: NodeId
+    keyed: bool = True
+    accept_children: bool = True
+    repr_: Optional[str] = None
+    data: Any = None
 
     def line_repr(self, depth: int, **kwargs: Any) -> Tuple[str, str]:
         """Control how node is displayed in tree representation.
-        _
-        ├── one                                           end
-        │   └── two                                     myEnd
-        └── three
+        First returned string is how node is represented on left, second string is how node is represented on right.
+
+        MyTree
+        ├── one                                        OneEnd
+        │   └── two                                    twoEnd
+        └── three                                    threeEnd
         """
-        if self.repr is not None:
-            return self.repr, ""
+        if self.repr_ is not None:
+            return self.repr_, ""
         if not self.accept_children:
-            return str(self.data), ""
+            if hasattr(self.data, "__str__"):
+                return str(self.data), ""
+            return "", ""
         if self.keyed:
             return "{}", ""
         return "[]", ""
 
-    def __eq__(self, other: Any) -> bool:
-        if not isinstance(other, self.__class__):
-            return False
-        return self.identifier == other.identifier
 
-    def __str__(self) -> str:
-        return "%s, id=%s" % (self.__class__.__name__, self.identifier)
+class AutoIdNode(Node):
+    def __init__(
+        self,
+        identifier: Optional[NodeId] = None,
+        keyed: bool = True,
+        accept_children: bool = True,
+        repr_: Optional[str] = None,
+        data: Any = None,
+    ):
+
+        self._auto_generated_id: bool
+        identifier_: NodeId
+
+        if identifier is None:
+            identifier_ = str(uuid.uuid4())
+            self._auto_generated_id = True
+        else:
+            identifier_ = identifier
+            self._auto_generated_id = False
 
-    def __repr__(self) -> str:
-        return self.__str__()
+        super(AutoIdNode, self).__init__(
+            identifier=identifier_,
+            keyed=keyed,
+            accept_children=accept_children,
+            repr_=repr_,
+            data=data,
+        )

lighttree/tree.py

@@ -11,6 +11,9 @@
     cast,
     Dict,
     Any,
+    Generic,
+    TypeVar,
+    Iterator,
 )
 from collections import defaultdict
 from operator import itemgetter
@@ -22,12 +25,14 @@
 
 # root has no key (None), keyed node has children with str keys, unkeyed node has children with int keys
 Key = Union[None, str, int]
-KeyedNode = Tuple[Key, Node]
 KeyedTree = Tuple[Key, "Tree"]
 Path = str
 
+GenericNode = TypeVar("GenericNode", bound=Node)
+KeyedNode = Tuple[Key, GenericNode]
 
-class Tree(object):
+
+class Tree(Generic[GenericNode]):
 
     """Principles:
     - each node is identified by an id
@@ -49,7 +54,7 @@ def __init__(self, path_separator: str = ".") -> None:
         # nodes references and hierarchy in tree
         self.root: Optional[NodeId] = None
         # node identifier -> node
-        self._nodes_map: Dict[NodeId, Node] = {}
+        self._nodes_map: Dict[NodeId, GenericNode] = {}
         # node identifier -> parent node identifier
         self._nodes_parent: Dict[NodeId, Optional[NodeId]] = defaultdict(lambda: None)
         # "map" node identifier -> map of children nodes identifier -> key
@@ -61,10 +66,6 @@ def __contains__(self, identifier: NodeId) -> bool:
         return identifier in self._nodes_map
 
     def get(self, nid: NodeId) -> KeyedNode:
-        """Get a node by its id.
-        :param nid: str, identifier of node to fetch
-        :rtype: lighttree.node.Node
-        """
         self._ensure_present(nid)
         return self.get_key(nid), self._nodes_map[nid]
 
@@ -128,13 +129,12 @@ def list(
         self,
         id_in: Optional[Sequence[NodeId]] = None,
         depth_in: Optional[Sequence[int]] = None,
-        filter_: Optional[Callable[[Node], bool]] = None,
+        filter_: Optional[Callable[[GenericNode], bool]] = None,
     ) -> 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
         :param filter\_: function, optional, filtering function to apply to each node
-        :rtype: list of lighttree.node.Node
         """
         return [
             (self.get_key(nid), node)
@@ -167,7 +167,7 @@ def _ensure_present(
             raise NotFoundNodeError("Node id <%s> doesn't exist in tree" % nid)
         return nid
 
-    def _validate_node_insertion(self, node: Node) -> None:
+    def _validate_node_insertion(self, node: GenericNode) -> None:
         if node.identifier in self._nodes_map.keys():
             raise DuplicatedNodeError(
                 "Can't create node with id '%s'" % node.identifier
@@ -330,24 +330,12 @@ def leaves_ids(self, nid: Optional[NodeId] = None) -> List[NodeId]:
 
     def insert(
         self,
-        item: Union[Node, "Tree"],
+        item: Union[GenericNode, "Tree"],
         parent_id: Optional[NodeId] = None,
         child_id: Optional[NodeId] = None,
         child_id_below: Optional[NodeId] = None,
         key: Key = None,
     ) -> "Tree":
-        if isinstance(item, Node):
-            if child_id_below is not None:
-                raise ValueError(
-                    '"child_id_below" parameter is reserved to Tree insertion.'
-                )
-            self.insert_node(
-                node=item,
-                parent_id=parent_id,
-                child_id=child_id,
-                key=key,
-            )
-            return self
         if isinstance(item, Tree):
             self.insert_tree(
                 new_tree=item,
@@ -357,13 +345,22 @@ def insert(
                 key=key,
             )
             return self
-        raise ValueError(
-            '"item" parameter must either be a Node, or a Tree, got <%s>.' % type(item)
+        # item is GenericNode
+        if child_id_below is not None:
+            raise ValueError(
+                '"child_id_below" parameter is reserved to Tree insertion.'
+            )
+        self.insert_node(
+            node=item,
+            parent_id=parent_id,
+            child_id=child_id,
+            key=key,
         )
+        return self
 
     def insert_node(
         self,
-        node: Node,
+        node: GenericNode,
         parent_id: Optional[NodeId] = None,
         child_id: Optional[NodeId] = None,
         key: Key = None,
@@ -386,7 +383,7 @@ def insert_node(
 
     def _insert_node_below(
         self,
-        node: Node,
+        node: GenericNode,
         parent_id: Optional[NodeId],
         key: Key,
     ) -> None:
@@ -433,7 +430,7 @@ def _insert_node_below(
         self._nodes_map[node_id] = node
         self._nodes_parent[node_id] = parent_id
 
-    def _insert_node_above(self, node: Node, child_id: NodeId, key: Key) -> None:
+    def _insert_node_above(self, node: GenericNode, child_id: NodeId, key: Key) -> None:
         self._ensure_present(child_id)
         # get parent_id before dropping subtree
         try:
@@ -601,10 +598,10 @@ def expand_tree(
         self,
         nid: Optional[NodeId] = None,
         mode: str = "depth",
-        filter_: Optional[Callable[[Union[None, str, int], Node], bool]] = None,
+        filter_: Optional[Callable[[Union[None, str, int], GenericNode], bool]] = None,
         filter_through: bool = False,
         reverse: bool = False,
-    ) -> Iterable[KeyedNode]:
+    ) -> Iterator[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
@@ -615,7 +612,6 @@ def expand_tree(
         :param filter_through: if True, excluded nodes don't exclude their children.
         :param reverse: the ``reverse`` param for sorting :class:`Node` objects in the same level
         :return: node ids that satisfy the conditions if ``id_only`` is True, else nodes.
-        :rtype: generator
         """
         if mode not in ("depth", "width"):
             raise NotImplementedError("Traversal mode '%s' is not supported" % mode)
@@ -656,7 +652,7 @@ def expand_tree(
     def show(
         self,
         nid: Optional[NodeId] = None,
-        filter_: Optional[Callable[[Node], bool]] = None,
+        filter_: Optional[Callable[[GenericNode], bool]] = None,
         display_key: bool = True,
         reverse: bool = False,
         line_type: str = "ascii-ex",
@@ -676,8 +672,6 @@ def show(
         :param limit: int, truncate tree display to this number of lines
         :param kwargs: kwargs params passed to node ``line_repr`` method
         :param line_max_length
-        :rtype: unicode in python2, str in python3
-
         """
         output = ""
 
@@ -714,10 +708,10 @@ def show(
     def _iter_nodes_with_location(
         self,
         nid: Optional[NodeId],
-        filter_: Optional[Callable[[Node], bool]],
+        filter_: Optional[Callable[[GenericNode], bool]],
         reverse: bool,
         is_last_list: Optional[List[bool]] = None,
-    ) -> Iterable[Tuple[Tuple[bool, ...], Key, Node]]:
+    ) -> Iterable[Tuple[Tuple[bool, ...], Key, GenericNode]]:
         """Yield nodes with information on how they are placed.
         :param nid: starting node identifier
         :param filter_: filter function applied on nodes