import textwrap
import contextvars
from pprint import isrecursive
from itertools import count
from collections import abc
_recursive_ids = contextvars.ContextVar('recursive')
[docs]
class TreePrinter:
"""Default printer for printree.
Uses unicode characters.
"""
ROOT = '┐'
EDGE = '│ '
BRANCH_NEXT = '├── '
BRANCH_LAST = '└── '
ARROW = '→'
def __init__(self, depth: int = None, annotated: bool = False):
"""
:param depth: If the data structure being printed is too deep, the next contained level is replaced by [...]. By default, there is no constraint on the depth of the objects being formatted.
:param annotated: Whether or not to include annotations for branches, like the object type and amount of children.
"""
self.level = 0
self.depth = depth
self.annotated = bool(annotated)
@property
def depth(self) -> int:
"""Maximum depth to traverse while creating the tree representation."""
return self._depth
@depth.setter
def depth(self, value):
if not (isinstance(value, int) or value is None):
raise TypeError(f"Expected depth to be an int or None. Got '{type(value).__name__}' instead.")
if isinstance(value, int) and value < 0:
raise ValueError(f"Depth must be a positive integer or zero. Got '{value}' instead.")
self._depth = value if value else float("inf")
def ptree(self, obj):
self.level = 0
def f():
_recursive_ids.set(set())
for i in _itree(obj, self, subscription=self.ROOT, depth=self.depth):
print(i)
contextvars.copy_context().run(f)
def ftree(self, obj):
self.level = 0
def f():
_recursive_ids.set(set())
return "\n".join(_itree(obj, self, subscription=self.ROOT, depth=self.depth))
return contextvars.copy_context().run(f)
[docs]
class AsciiPrinter(TreePrinter):
"""A printer that uses ASCII characters only."""
ROOT = '.'
EDGE = '| '
BRANCH_NEXT = '|-- '
BRANCH_LAST = '`-- '
ARROW = '->'
[docs]
def ptree(obj, depth: int = None, annotated: bool = False) -> None:
"""Print a tree-like representation of the given object data structure.
:py:class:`collections.abc.Iterable` instances will be branches, except for :py:class:`str` and :py:class:`bytes`.
All other objects will be leaves.
:param depth: If the data structure being printed is too deep, the next contained level is replaced by [...]. By default, there is no constraint on the depth of the objects being formatted.
:param annotated: Whether to include annotations for branches, like the object type and amount of children.
Examples:
>>> dct = {"A": {"x\\ny", (42, -17, 0.01), True}, "B": 42}
>>> ptree(dct)
┐
├─ A
│ ├─ 0: x
│ │ y
│ ├─ 1
│ │ ├─ 0: 42
│ │ ├─ 1: -17
│ │ └─ 2: 0.01
│ └─ 2: True
└─ B: 42
>>> ptree(dct, annotated=True, depth=2)
┐ → dict[items=2]
├─ A → set[items=3]
│ ├─ 0: x
│ │ y
│ ├─ 1 → tuple[items=3] [...]
│ └─ 2: True
└─ B: 42
"""
TreePrinter(depth=depth, annotated=annotated).ptree(obj)
[docs]
def ftree(obj, depth:int=None, annotated:bool=False) -> str:
"""Return the formatted tree representation of the given object data structure as a string. Arguments are same as `ptree`"""
return TreePrinter(depth=depth, annotated=annotated).ftree(obj)
def _newline_repr(obj_repr, prefix) -> str:
counter = count()
newline = lambda x: next(counter) != 0
return textwrap.indent(obj_repr, prefix, newline)
def _itree(obj, formatter, subscription, prefix="", last=False, level=0, depth=0):
formatter.level = level
sprout = level > 0
children = []
objid = id(obj)
recursive = isrecursive(obj)
recursive_ids = _recursive_ids.get()
sprout_repr = ': ' if sprout else ''
newlevel = ' ' if last else formatter.EDGE
newline_prefix = f"{prefix}{newlevel}"
newprefix = f"{prefix}{formatter.BRANCH_LAST if last else formatter.BRANCH_NEXT}" if sprout else ""
subscription_repr = f'{newprefix}{_newline_repr(f"{subscription}", newline_prefix)}'
if recursive and objid in recursive_ids:
item_repr = f"{sprout_repr}<Recursion on {type(obj).__name__} with id={objid}>"
elif isinstance(obj, (str, bytes)):
# Indent new lines with a prefix so that a string like "new\nline" adjusts to:
# ...
# |- 42: new
# | line
# ...
# for this, calculate how many characters each new line should have for padding
prefix_len = len(prefix) # how much we have to copy before subscription string
last_line = subscription_repr.expandtabs().splitlines()[-1]
newline_padding = len(last_line[prefix_len:]) + prefix_len + 2 # last 2 are ": "
item_repr = _newline_repr(f"{sprout_repr}{obj}", f"{last_line[:prefix_len] + newlevel:<{newline_padding}}")
elif isinstance(obj, abc.Iterable):
# for other iterable objects, enumerate to track subscription and child count
ismap = isinstance(obj, abc.Mapping)
enumerateable = obj.items() if ismap else obj
accessor = (lambda i, v: (i, *v)) if ismap else lambda i, v: (i, i, v)
enumerated = enumerate(enumerateable)
children.extend(accessor(*enum) for enum in enumerated)
contents = f'items={len(children)}' if children else "empty"
item_repr = f' {formatter.ARROW} {type(obj).__name__}[{contents}]' if formatter.annotated else ''
if children and level == depth:
item_repr = f"{item_repr} [...]"
children.clear() # avoid deeper traversal
else:
item_repr = f"{sprout_repr}{obj}"
if recursive:
recursive_ids.add(objid)
yield f"{subscription_repr}{item_repr}"
child_count = len(children)
prefix += newlevel if sprout else "" # only add level prefix starting at level 1
for index, key, value in children:
yield from _itree(value, formatter, key, prefix, index == (child_count - 1), level + 1, depth)