From 49565454d8cb6bd3819d025141ca4c9d80867550 Mon Sep 17 00:00:00 2001 From: Daniel Baumann Date: Wed, 5 Mar 2025 09:29:32 +0100 Subject: [PATCH] Merging upstream version 1.7.1. Signed-off-by: Daniel Baumann --- .github/workflows/pypi-publish.yml | 4 +- .github/workflows/python-package.yml | 33 +++- .travis.yml | 15 -- setup.py | 7 +- treelib/node.py | 87 ++++++---- treelib/py.typed | 0 treelib/tree.py | 237 +++++++++++++++++---------- 7 files changed, 235 insertions(+), 148 deletions(-) delete mode 100644 .travis.yml create mode 100644 treelib/py.typed diff --git a/.github/workflows/pypi-publish.yml b/.github/workflows/pypi-publish.yml index bdaab28..50795dc 100644 --- a/.github/workflows/pypi-publish.yml +++ b/.github/workflows/pypi-publish.yml @@ -21,9 +21,9 @@ jobs: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v4 - name: Set up Python - uses: actions/setup-python@v3 + uses: actions/setup-python@v5.4.0 with: python-version: '3.x' - name: Install dependencies diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml index fba950c..ae64e35 100644 --- a/.github/workflows/python-package.yml +++ b/.github/workflows/python-package.yml @@ -6,21 +6,38 @@ on: push: branches: [ "master", "dev" ] pull_request: - branches: [ "master" ] + types: [opened, synchronize, reopened] jobs: - build: - - runs-on: ubuntu-latest + build-compat: + runs-on: ubuntu-20.04 strategy: fail-fast: false matrix: - python-version: ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12.0-rc.1"] - + python-version: ["3.6", "3.7"] steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v3 + uses: actions/setup-python@v5.4.0 + with: + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: | + if [ -f requirements.txt ]; then pip install -r requirements.txt; fi + if [ -f requirements-t-pre37.txt ]; then pip install -r requirements-t-pre37.txt; fi + - name: Test with nosetest + run: | + ./scripts/testing.sh + build-latest: + runs-on: ubuntu-22.04 + strategy: + fail-fast: false + matrix: + python-version: ["3.8", "3.9", "3.10", "3.11", "3.12", "3.13"] + steps: + - uses: actions/checkout@v4 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5.4.0 with: python-version: ${{ matrix.python-version }} - name: Install dependencies diff --git a/.travis.yml b/.travis.yml deleted file mode 100644 index 944c475..0000000 --- a/.travis.yml +++ /dev/null @@ -1,15 +0,0 @@ -language: python -python: - - "2.7" - - "3.4" - - "3.5" - - "3.6" - - "3.7" - - "3.8" - -install: - - "python setup.py install" - -script: - - pip install -r requirements-testing.txt - - sh testing.sh diff --git a/setup.py b/setup.py index 9331bec..747574a 100644 --- a/setup.py +++ b/setup.py @@ -1,6 +1,6 @@ from setuptools import setup -__version__ = "1.7.0" +__version__ = "1.7.1" setup( @@ -15,6 +15,8 @@ setup( packages=["treelib"], keywords=["data structure", "tree", "tools"], install_requires=["six"], + include_package_data=True, + package_data={"": ["py.typed"]}, classifiers=[ "Development Status :: 5 - Production/Stable", "Environment :: Console", @@ -23,8 +25,6 @@ setup( "Operating System :: OS Independent", "Programming Language :: Python", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.4", - "Programming Language :: Python :: 3.5", "Programming Language :: Python :: 3.6", "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", @@ -32,6 +32,7 @@ setup( "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", "Topic :: Software Development :: Libraries :: Python Modules", ], ) diff --git a/treelib/node.py b/treelib/node.py index 958c734..d0e789b 100644 --- a/treelib/node.py +++ b/treelib/node.py @@ -27,12 +27,19 @@ from __future__ import unicode_literals import copy import uuid +import sys from collections import defaultdict from warnings import warn +from typing import cast, List, Any, Optional, Union from .exceptions import NodePropertyError from .misc import deprecated +if sys.version_info >= (3, 9): + StrList = list[str] +else: + StrList = List[str] # Python 3.8 and earlier + class Node(object): """ @@ -43,11 +50,17 @@ class Node(object): #: Mode constants for routine `update_fpointer()`. (ADD, DELETE, INSERT, REPLACE) = list(range(4)) - def __init__(self, tag=None, identifier=None, expanded=True, data=None): + def __init__( + self, + tag: Optional[str] = None, + identifier: Optional[str] = None, + expanded: bool = True, + data: Any = None, + ) -> None: """Create a new Node object to be placed inside a Tree object""" #: if given as a parameter, must be unique - self._identifier = None + self._identifier: Optional[str] = None self._set_identifier(identifier) #: None or something else @@ -61,24 +74,24 @@ class Node(object): self.expanded = expanded #: identifier of the parent's node : - self._predecessor = {} + self._predecessor: dict = {} #: identifier(s) of the soons' node(s) : - self._successors = defaultdict(list) + self._successors: dict = defaultdict(list) #: User payload associated with this node. self.data = data # for retro-compatibility on bpointer/fpointer - self._initial_tree_id = None + self._initial_tree_id: Optional[str] = None - def __lt__(self, other): + def __lt__(self, other) -> bool: return self.tag < other.tag - def set_initial_tree_id(self, tree_id): + def set_initial_tree_id(self, tree_id: str) -> None: if self._initial_tree_id is None: self._initial_tree_id = tree_id - def _set_identifier(self, nid): + def _set_identifier(self, nid: Optional[str]) -> None: """Initialize self._set_identifier""" if nid is None: self._identifier = str(uuid.uuid1()) @@ -98,11 +111,11 @@ class Node(object): @bpointer.setter @deprecated(alias="node.set_predecessor") - def bpointer(self, value): + def bpointer(self, value) -> None: self.set_predecessor(value, self._initial_tree_id) @deprecated(alias="node.set_predecessor") - def update_bpointer(self, nid): + def update_bpointer(self, nid) -> None: self.set_predecessor(nid, self._initial_tree_id) @property @@ -118,7 +131,7 @@ class Node(object): @fpointer.setter @deprecated(alias="node.update_successors") - def fpointer(self, value): + def fpointer(self, value: Union[None, list, dict, set]) -> None: self.set_successors(value, tree_id=self._initial_tree_id) @deprecated(alias="node.update_successors") @@ -132,11 +145,11 @@ class Node(object): """ return self._predecessor[tree_id] - def set_predecessor(self, nid, tree_id): + def set_predecessor(self, nid: Optional[str], tree_id: Optional[str]) -> None: """Set the value of `_predecessor`.""" self._predecessor[tree_id] = nid - def successors(self, tree_id): + def successors(self, tree_id: Optional[str]) -> StrList: """ With a getting operator, a list of IDs of node's children is obtained. With a setting operator, the value can be list, set, or dict. For list or set, @@ -145,7 +158,9 @@ class Node(object): """ return self._successors[tree_id] - def set_successors(self, value, tree_id=None): + def set_successors( + self, value: Union[None, list, dict, set], tree_id: Optional[str] = None + ) -> None: """Set the value of `_successors`.""" setter_lookup = { "NoneType": lambda x: list(), @@ -161,7 +176,13 @@ class Node(object): else: raise NotImplementedError("Unsupported value type %s" % t) - def update_successors(self, nid, mode=ADD, replace=None, tree_id=None): + def update_successors( + self, + nid: Optional[str], + mode: int = ADD, + replace: Optional[str] = None, + tree_id: Optional[str] = None, + ) -> None: """ Update the children list with different modes: addition (Node.ADD or Node.INSERT) and deletion (Node.DELETE). @@ -169,23 +190,23 @@ class Node(object): if nid is None: return - def _manipulator_append(): + def _manipulator_append() -> None: self.successors(tree_id).append(nid) - def _manipulator_delete(): + def _manipulator_delete() -> None: if nid in self.successors(tree_id): self.successors(tree_id).remove(nid) else: warn("Nid %s wasn't present in fpointer" % nid) - def _manipulator_insert(): + def _manipulator_insert() -> None: warn("WARNING: INSERT is deprecated to ADD mode") self.update_successors(nid, tree_id=tree_id) - def _manipulator_replace(): + def _manipulator_replace() -> None: if replace is None: raise NodePropertyError( - 'Argument "repalce" should be provided when mode is {}'.format(mode) + 'Argument "replace" should be provided when mode is {}'.format(mode) ) ind = self.successors(tree_id).index(nid) self.successors(tree_id)[ind] = replace @@ -200,38 +221,38 @@ class Node(object): if mode not in manipulator_lookup: raise NotImplementedError("Unsupported node updating mode %s" % str(mode)) - f_name = manipulator_lookup.get(mode) + f_name = cast(str, manipulator_lookup.get(mode)) f = locals()[f_name] return f() @property - def identifier(self): + def identifier(self) -> str: """ The unique ID of a node within the scope of a tree. This attribute can be accessed and modified with ``.`` and ``=`` operator respectively. """ - return self._identifier + return cast(str, self._identifier) - def clone_pointers(self, former_tree_id, new_tree_id): + def clone_pointers(self, former_tree_id: str, new_tree_id: str) -> None: former_bpointer = self.predecessor(former_tree_id) self.set_predecessor(former_bpointer, new_tree_id) former_fpointer = self.successors(former_tree_id) # fpointer is a list and would be copied by reference without deepcopy self.set_successors(copy.deepcopy(former_fpointer), tree_id=new_tree_id) - def reset_pointers(self, tree_id): + def reset_pointers(self, tree_id) -> None: self.set_predecessor(None, tree_id) self.set_successors([], tree_id=tree_id) - @identifier.setter - def identifier(self, value): + @identifier.setter # type: ignore + def identifier(self, value: str) -> None: """Set the value of `_identifier`.""" if value is None: print("WARNING: node ID can not be None") else: self._set_identifier(value) - def is_leaf(self, tree_id=None): + def is_leaf(self, tree_id: Optional[str] = None) -> bool: """Return true if current node has no children.""" if tree_id is None: # for retro-compatilibity @@ -245,7 +266,7 @@ class Node(object): else: return False - def is_root(self, tree_id=None): + def is_root(self, tree_id: Optional[str] = None) -> bool: """Return true if self has no parent, i.e. as root.""" if tree_id is None: # for retro-compatilibity @@ -257,19 +278,19 @@ class Node(object): return self.predecessor(tree_id) is None @property - def tag(self): + def tag(self) -> str: """ The readable node name for human. This attribute can be accessed and modified with ``.`` and ``=`` operator respectively. """ - return self._tag + return cast(str, self._tag) @tag.setter - def tag(self, value): + def tag(self, value: Optional[str]) -> None: """Set the value of `_tag`.""" self._tag = value if value is not None else None - def __repr__(self): + def __repr__(self) -> str: name = self.__class__.__name__ kwargs = [ "tag={0}".format(self.tag), diff --git a/treelib/py.typed b/treelib/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/treelib/tree.py b/treelib/tree.py index 1cc9ac8..d9b761f 100644 --- a/treelib/tree.py +++ b/treelib/tree.py @@ -32,16 +32,18 @@ from __future__ import unicode_literals try: from builtins import str as text except ImportError: - from __builtin__ import str as text + from __builtin__ import str as text # type: ignore import codecs import json import uuid +import sys from copy import deepcopy from six import python_2_unicode_compatible, iteritems +from typing import cast, List, Any, Callable, Optional, Union try: - from StringIO import StringIO + from StringIO import StringIO # type: ignore except ImportError: from io import StringIO @@ -55,6 +57,16 @@ from .exceptions import ( ) from .node import Node +if sys.version_info >= (3, 9): + StrList = list[str] + StrLList = list[list[str]] + NodeList = list[Node] +else: + StrList = List[str] # Python 3.8 and earlier + StrLList = List[List[str]] + NodeList = List[Node] + + __author__ = "chenxm" @@ -66,14 +78,20 @@ class Tree(object): (ROOT, DEPTH, WIDTH, ZIGZAG) = list(range(4)) node_class = Node - def __contains__(self, identifier): + def __contains__(self, identifier: str) -> bool: return identifier in self.nodes.keys() - def __init__(self, tree=None, deep=False, node_class=None, identifier=None): + def __init__( + self, + tree=None, + deep: bool = False, + node_class=None, + identifier: Optional[str] = None, + ) -> None: """Initiate a new tree or copy another tree with a shallow or deep copy. """ - self._identifier = None + self._identifier: Optional[str] = None self._set_identifier(identifier) if node_class: @@ -85,7 +103,7 @@ class Tree(object): #: Get or set the identifier of the root. This attribute can be accessed and modified #: with ``.`` and ``=`` operator respectively. - self.root = None + self.root: Optional[str] = None if tree is not None: self.root = tree.root @@ -95,7 +113,12 @@ class Tree(object): if tree.identifier != self._identifier: new_node.clone_pointers(tree.identifier, self._identifier) - def _clone(self, identifier=None, with_tree=False, deep=False): + def _clone( + self, + identifier: Optional[str] = None, + with_tree: bool = False, + deep: bool = False, + ): """Clone current instance, with or without tree. Method intended to be overloaded, to avoid rewriting whole "subtree" and "remove_subtree" methods when @@ -122,28 +145,28 @@ class Tree(object): ) @property - def identifier(self): + def identifier(self) -> Optional[str]: return self._identifier - def _set_identifier(self, nid): + def _set_identifier(self, nid: Optional[str]) -> None: """Initialize self._set_identifier""" if nid is None: self._identifier = str(uuid.uuid1()) else: self._identifier = nid - def __getitem__(self, key): + def __getitem__(self, key: str) -> Node: """Return _nodes[key]""" try: return self._nodes[key] except KeyError: raise NodeIDAbsentError("Node '%s' is not in the tree" % key) - def __len__(self): + def __len__(self) -> int: """Return len(_nodes)""" return len(self._nodes) - def __str__(self): + def __str__(self) -> str: self._reader = "" def write(line): @@ -154,16 +177,16 @@ class Tree(object): def __print_backend( self, - nid=None, - level=ROOT, - idhidden=True, - filter=None, - key=None, - reverse=False, + nid: Optional[str] = None, + level: int = ROOT, + idhidden: bool = True, + filter: Optional[Callable[[Node], bool]] = None, + key: Optional[Callable[[Node], Node]] = None, + reverse: bool = False, line_type="ascii-ex", data_property=None, - sorting=True, - func=print, + sorting: bool = True, + func: Callable[[bytes], None] = print, ): """ Another implementation of printing tree using Stack @@ -193,12 +216,12 @@ class Tree(object): if data_property: if idhidden: - def get_label(node): + def get_label(node: Node) -> str: return getattr(node.data, data_property) else: - def get_label(node): + def get_label(node: Node) -> str: return "%s[%s]" % ( getattr(node.data, data_property), node.identifier, @@ -207,12 +230,12 @@ class Tree(object): else: if idhidden: - def get_label(node): + def get_label(node: Node) -> str: return node.tag else: - def get_label(node): + def get_label(node: Node) -> str: return "%s[%s]" % (node.tag, node.identifier) # legacy ordering @@ -228,7 +251,16 @@ class Tree(object): label = get_label(node) func("{0}{1}".format(pre, label).encode("utf-8")) - def __get(self, nid, level, filter_, key, reverse, line_type, sorting): + def __get( + self, + nid: Optional[str], + level: int, + filter_: Optional[Callable[[Node], bool]], + key: Optional[Callable[[Node], Node]], + reverse: bool, + line_type: str, + sorting: bool, + ): # default filter if filter_ is None: @@ -247,10 +279,20 @@ class Tree(object): return self.__get_iter(nid, level, filter_, key, reverse, dt, [], sorting) - def __get_iter(self, nid, level, filter_, key, reverse, dt, is_last, sorting): + def __get_iter( + self, + nid: Optional[str], + level: int, + filter_: Callable[[Node], bool], + key: Optional[Callable[[Node], Node]], + reverse: bool, + dt, # tuple[str, str, str] + is_last: list, + sorting: bool, + ): dt_vertical_line, dt_line_box, dt_line_corner = dt - nid = self.root if nid is None else nid + nid = cast(str, self.root if nid is None else nid) node = self[nid] if level == self.ROOT: @@ -274,7 +316,7 @@ class Tree(object): if key: children.sort(key=key, reverse=reverse) elif reverse: - children = reversed(children) + children = list(reversed(children)) level += 1 for idx, child in enumerate(children): is_last.append(idx == idxlast) @@ -288,13 +330,13 @@ class Tree(object): """set self[nid].bpointer""" self[nid].set_predecessor(parent_id, self._identifier) - def __update_fpointer(self, nid, child_id, mode): + def __update_fpointer(self, nid: str, child_id: str, mode: int): if nid is None: return else: self[nid].update_successors(child_id, mode, tree_id=self._identifier) - def add_node(self, node, parent=None): + def add_node(self, node: Node, parent: Optional[Union[Node, str]] = None) -> None: """ Add a new node object to the tree and make the parent as the root by default. @@ -320,16 +362,17 @@ class Tree(object): elif not self.contains(pid): raise NodeIDAbsentError("Parent node '%s' " "is not in the tree" % pid) + pid = cast(str, pid) self._nodes.update({node.identifier: node}) self.__update_fpointer(pid, node.identifier, self.node_class.ADD) self.__update_bpointer(node.identifier, pid) - node.set_initial_tree_id(self._identifier) + node.set_initial_tree_id(cast(str, self._identifier)) - def all_nodes(self): + def all_nodes(self) -> NodeList: """Return all nodes in a list""" return list(self._nodes.values()) - def all_nodes_itr(self): + def all_nodes_itr(self) -> Any: """ Returns all nodes in an iterator. Added by William Rusnack @@ -368,7 +411,7 @@ class Tree(object): ascendant_level = self.level(ascendant) return None - def children(self, nid): + def children(self, nid: str) -> NodeList: """ Return the children (Node) list of nid. Empty list is returned if nid does not exist @@ -379,7 +422,13 @@ class Tree(object): """Check if the tree contains node of given id""" return True if nid in self._nodes else False - def create_node(self, tag=None, identifier=None, parent=None, data=None): + def create_node( + self, + tag: Optional[str] = None, + identifier: Optional[str] = None, + parent: Optional[Union[Node, str]] = None, + data: Any = None, + ) -> Node: """ Create a child node for given @parent node. If ``identifier`` is absent, a UUID will be generated automatically. @@ -388,7 +437,7 @@ class Tree(object): self.add_node(node, parent) return node - def depth(self, node=None): + def depth(self, node: Optional[Node] = None) -> int: """ Get the maximum level of this tree or the level of the given node. @@ -415,7 +464,13 @@ class Tree(object): return ret def expand_tree( - self, nid=None, mode=DEPTH, filter=None, key=None, reverse=False, sorting=True + self, + nid: Optional[str] = None, + mode: int = DEPTH, + filter: Optional[Callable[[Node], bool]] = None, + key: Optional[Callable[[Node], Node]] = None, + reverse: bool = False, + sorting: bool = True, ): """ Python generator to traverse the tree (or a subtree) with optional @@ -439,7 +494,7 @@ class Tree(object): :return: Node IDs that satisfy the conditions :rtype: generator object """ - nid = self.root if nid is None else nid + nid = cast(str, self.root if nid is None else nid) if not self.contains(nid): raise NodeIDAbsentError("Node '%s' is not in the tree" % nid) @@ -470,7 +525,7 @@ class Tree(object): elif mode is self.ZIGZAG: # Suggested by Ilya Kuprik (ilya-spy@ynadex.ru). - stack_fw = [] + stack_fw: NodeList = [] queue.reverse() stack = stack_bw = queue direction = False @@ -493,7 +548,7 @@ class Tree(object): else: raise ValueError("Traversal mode '{}' is not supported".format(mode)) - def filter_nodes(self, func): + def filter_nodes(self, func: Callable[[Node], bool]): """ Filters all nodes by function. @@ -504,7 +559,7 @@ class Tree(object): """ return filter(func, self.all_nodes_itr()) - def get_node(self, nid): + def get_node(self, nid: Optional[str]) -> Optional[Node]: """ Get the object of the node with ID of ``nid``. @@ -531,7 +586,7 @@ class Tree(object): fpointer = [] return fpointer - def leaves(self, nid=None): + def leaves(self, nid: Optional[str] = None) -> NodeList: """Get leaves of the whole tree or a subtree.""" leaves = [] if nid is None: @@ -555,7 +610,7 @@ class Tree(object): """ return len([n for n in self.rsearch(nid, filter)]) - 1 - def link_past_node(self, nid): + def link_past_node(self, nid: str): """ Delete a node by linking past it. @@ -617,7 +672,7 @@ class Tree(object): """Return a dict form of nodes in a tree: {id: node_instance}.""" return self._nodes - def parent(self, nid): + def parent(self, nid: str) -> Optional[Node]: """Get parent :class:`Node` object of given id.""" if not self.contains(nid): raise NodeIDAbsentError("Node '%s' is not in the tree" % nid) @@ -628,7 +683,7 @@ class Tree(object): return self[pid] - def merge(self, nid, new_tree, deep=False): + def merge(self, nid: str, new_tree, deep: bool = False): """Patch @new_tree on current tree by pasting new_tree root children on current tree @nid node. Consider the following tree: @@ -667,7 +722,7 @@ class Tree(object): for child in new_tree.children(new_tree.root): self.paste(nid=nid, new_tree=new_tree.subtree(child.identifier), deep=deep) - def paste(self, nid, new_tree, deep=False): + def paste(self, nid: str, new_tree, deep: bool = False): """ Paste a @new_tree to the original one by linking the root of new tree to given node (nid). @@ -698,7 +753,7 @@ class Tree(object): self.__update_bpointer(new_tree.root, nid) self.__update_fpointer(nid, new_tree.root, self.node_class.ADD) - def paths_to_leaves(self): + def paths_to_leaves(self) -> StrLList: """ Use this function to get the identifiers allowing to go from the root nodes to each leaf. @@ -735,7 +790,7 @@ class Tree(object): return res - def remove_node(self, identifier): + def remove_node(self, identifier: str) -> int: """Remove a node indicated by 'identifier' with all its successors. Return the number of removed nodes. """ @@ -762,7 +817,7 @@ class Tree(object): self.nodes.pop(id_) return len(removed) - def remove_subtree(self, nid, identifier=None): + def remove_subtree(self, nid: str, identifier: Optional[str] = None): """ Get a subtree with ``nid`` being the root. If nid is None, an empty tree is returned. @@ -798,14 +853,16 @@ class Tree(object): if id_ == self.root: self.root = None st._nodes.update({id_: self._nodes.pop(id_)}) - st[id_].clone_pointers(self._identifier, st.identifier) + st[id_].clone_pointers( + cast(str, self._identifier), cast(str, st.identifier) + ) st[id_].reset_pointers(self._identifier) if id_ == nid: st[id_].set_predecessor(None, st.identifier) self.__update_fpointer(parent, nid, self.node_class.DELETE) return st - def rsearch(self, nid, filter=None): + def rsearch(self, nid: str, filter: Optional[Callable[[Node], bool]] = None): """ Traverse the tree branch along the branch from nid to its ancestors (until root). @@ -833,16 +890,16 @@ class Tree(object): def save2file( self, - filename, - nid=None, - level=ROOT, - idhidden=True, - filter=None, - key=None, - reverse=False, - line_type="ascii-ex", - data_property=None, - sorting=True, + filename: str, + nid: Optional[str] = None, + level: int = ROOT, + idhidden: bool = True, + filter: Optional[Callable[[Node], bool]] = None, + key: Optional[Callable[[Node], Node]] = None, + reverse: bool = False, + line_type: str = "ascii-ex", + data_property: Optional[str] = None, + sorting: bool = True, ): """ Save the tree into file for offline analysis. @@ -869,16 +926,16 @@ class Tree(object): def show( self, - nid=None, - level=ROOT, - idhidden=True, - filter=None, - key=None, - reverse=False, - line_type="ascii-ex", - data_property=None, - stdout=True, - sorting=True, + nid: Optional[str] = None, + level: int = ROOT, + idhidden: bool = True, + filter: Optional[Callable[[Node], bool]] = None, + key: Optional[Callable[[Node], Node]] = None, + reverse: bool = False, + line_type: str = "ascii-ex", + data_property: Optional[str] = None, + stdout: bool = True, + sorting: bool = True, ): """ Print the tree structure in hierarchy style. @@ -902,6 +959,7 @@ class Tree(object): :param reverse: the ``reverse`` param for sorting :class:`Node` objects in the same level. :param line_type: :param data_property: the property on the node data object to be printed. + :param stdout: if True print it, if False return printing. :param sorting: if True perform node sorting, if False return nodes in original insertion order. In latter case @key and @reverse parameters are ignored. @@ -929,11 +987,11 @@ class Tree(object): print("Tree is empty") if stdout: - print(self._reader.encode("utf-8")) + print(self._reader) else: return self._reader - def siblings(self, nid): + def siblings(self, nid: str) -> NodeList: """ Return the siblings of given @nid. @@ -949,7 +1007,7 @@ class Tree(object): return siblings - def size(self, level=None): + def size(self, level: Optional[int] = None) -> int: """ Get the number of nodes of the whole tree if @level is not given. Otherwise, the total number of nodes at specific level @@ -977,7 +1035,7 @@ class Tree(object): "level should be an integer instead of '%s'" % type(level) ) - def subtree(self, nid, identifier=None): + def subtree(self, nid: str, identifier: Optional[str] = None): """ Return a shallow COPY of subtree with nid being the new root. If nid is None, return an empty tree. @@ -1002,13 +1060,15 @@ class Tree(object): st._nodes.update({self[node_n].identifier: self[node_n]}) # define nodes parent/children in this tree # all pointers are the same as copied tree, except the root - st[node_n].clone_pointers(self._identifier, st.identifier) + st[node_n].clone_pointers( + cast(str, self._identifier), cast(str, st.identifier) + ) if node_n == nid: # reset root parent for the new tree st[node_n].set_predecessor(None, st.identifier) return st - def update_node(self, nid, **attrs): + def update_node(self, nid: str, **attrs) -> None: """ Update node's attributes. @@ -1072,19 +1132,21 @@ class Tree(object): ) return tree_dict - def to_json(self, with_data=False, sort=True, reverse=False): + def to_json( + self, with_data: bool = False, sort: bool = True, reverse: bool = False + ): """To format the tree in JSON format.""" return json.dumps(self.to_dict(with_data=with_data, sort=sort, reverse=reverse)) def to_graphviz( self, - filename=None, - shape="circle", - graph="digraph", + filename: Optional[str] = None, + shape: str = "circle", + graph: str = "digraph", filter=None, key=None, - reverse=False, - sorting=True, + reverse: bool = False, + sorting: bool = True, ): """Exports the tree in the dot format of the graphviz software""" nodes, connections = [], [] @@ -1097,7 +1159,8 @@ class Tree(object): sorting=sorting, ): nid = self[n].identifier - state = '"{0}" [label="{1}", shape={2}]'.format(nid, self[n].tag, shape) + label = str(self[n].tag).translate(str.maketrans({'"': r"\""})) + state = '"{0}" [label="{1}", shape={2}]'.format(nid, label, shape) nodes.append(state) for c in self.children(nid): @@ -1108,7 +1171,7 @@ class Tree(object): # write nodes and connections to dot format is_plain_file = filename is not None if is_plain_file: - f = codecs.open(filename, "w", "utf-8") + f = codecs.open(cast(str, filename), "w", "utf-8") else: f = StringIO() @@ -1119,8 +1182,8 @@ class Tree(object): if len(connections) > 0: f.write("\n") - for c in connections: - f.write("\t" + c + "\n") + for cns in connections: + f.write("\t" + cns + "\n") f.write("}")