Skip to content

Commit

Permalink
recursion and tab fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
chrizzFTD committed Sep 7, 2020
1 parent 6701b0f commit de3c0b7
Show file tree
Hide file tree
Showing 3 changed files with 84 additions and 42 deletions.
34 changes: 16 additions & 18 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -22,10 +22,10 @@ Similar to pretty print ([pprint](https://docs.python.org/3/library/pprint.html)
```

Instances of [abc.Iterable](https://docs.python.org/3/library/collections.abc.html#collections.abc.Iterable) (with the exception of [str](https://docs.python.org/3/library/stdtypes.html#text-sequence-type-str) & [bytes](https://docs.python.org/3/library/stdtypes.html#bytes-objects)) should be translated into a tree-like form.
Other objects will be considered "leaf nodes":
All other objects will be considered "leaf nodes":
```python
>>> dct = {
... "multi\nlined\n\ttabbed key": 1,
... "foo": [],
... True: {
... "uno": {"ABC", "XYZ"},
... "dos": r"B:\newline\tab\like.ext",
Expand All @@ -34,16 +34,13 @@ Other objects will be considered "leaf nodes":
... "numbers": (42, -17, 0.01)
... },
... },
... "foo": [],
... ("unsortable", ("tuple", "as", "key")):
... ["multi\nline\nfirst", "multi\nline\nlast"]
... {"multi\nlined\n\ttabbed key": "multi\nline\n\ttabbed value"}
... }
>>> dct['recursive_reference'] = dct
>>> dct["recursion"] = [1, dct, 2]
>>> ptree(dct)
`- . [items=5]
|- multi
| lined
| tabbed key: 1
`- . [items=4]
|- foo [empty]
|- True [items=3]
| |- dos: B:\newline\tab\like.ext
| |- tres [items=2]
Expand All @@ -55,13 +52,14 @@ Other objects will be considered "leaf nodes":
| `- uno [items=2]
| |- 0: ABC
| `- 1: XYZ
|- foo [empty]
|- ('unsortable', ('tuple', 'as', 'key')) [items=2]
| |- 0: multi
| | line
| | first
| `- 1: multi
| line
| last
`- recursive_reference: <Recursion on dict with id=140712966998864>
|- ('unsortable', ('tuple', 'as', 'key')) [items=1]
| `- multi
| lined
| tabbed key: multi
| line
| tabbed value
`- recursion [items=3]
|- 0: 1
|- 1: <Recursion on dict with id=2317960566912>
`- 2: 2
```
39 changes: 26 additions & 13 deletions printree/_ptree.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
import textwrap
import contextvars
from pprint import isrecursive
from itertools import count
from collections import abc

_root = contextvars.ContextVar('root')
_recursive_ids = contextvars.ContextVar('recursive')


def ptree(obj, /) -> None:
Expand Down Expand Up @@ -35,7 +36,7 @@ def ptree(obj, /) -> None:
`- C: <Recursion on dict with id=140712966998864>
"""
def f():
_root.set(obj)
_recursive_ids.set(set())
for i in _itree(obj):
print(i)
ctx = contextvars.copy_context()
Expand All @@ -45,7 +46,7 @@ def f():
def ftree(obj, /) -> str:
"""Return the formatted tree representation of the given object data structure as a string."""
def f():
_root.set(obj)
_recursive_ids.set(set())
return "\n".join(_itree(obj))
ctx = contextvars.copy_context()
return ctx.run(f)
Expand All @@ -59,18 +60,27 @@ def _newline_repr(obj_repr, /, prefix) -> str:

def _itree(obj, /, subscription=".", prefix="", last=True):
children = []
item_repr = f': {obj}'
if _root.get() is obj and subscription != ".": # recursive reference in container
item_repr = f": <Recursion on {type(obj).__name__} with id={id(object)}>"
level_suffix = ' ' if last else '| '
newline_prefix = f"{prefix}{level_suffix}"
subscription_repr = _newline_repr(f"{subscription}", newline_prefix)
recursive_ids = _recursive_ids.get()
recursive = isrecursive(obj)
objid = id(obj)
if recursive and objid in recursive_ids:
item_repr = f": <Recursion on {type(obj).__name__} with id={objid}>"
elif isinstance(obj, (str, bytes)):
# for string and bytes, indent new lines with an appropiate prefix so that
# a string line "new\nline" is adjusted to something like:
# for text, indent new lines with an appropiate prefix so that
# a string like "new\nline" is adjusted to something like:
# ...
# |- 42: new
# | line
# ...
newline_item_prefix = f'{prefix}{" " if last else "| "}{" " * len(f"{subscription}")}'
item_repr = _newline_repr(item_repr, newline_item_prefix)
# for this, calculate how many characters each new line should have for padding
# based on the last line from the subscription repr
prefix_len = len(newline_prefix)
no_prefix = subscription_repr.splitlines()[-1].expandtabs()[prefix_len:]
newline_padding = len(no_prefix) + prefix_len + 2 # last 2 are ": " below
item_repr = _newline_repr(f': {obj}', f"{newline_prefix:<{newline_padding}}")
elif isinstance(obj, abc.Iterable):
# for other iterable objects, sort and ennumerate so that we can anticipate what
# prefix we should use (e.g. are we the last item in the iteration?)
Expand All @@ -83,11 +93,14 @@ def _itree(obj, /, subscription=".", prefix="", last=True):
enumerated = enumerate(enumerateable)
children.extend(accessor(*enum) for enum in enumerated)
item_repr = f' [{items=}]' if (items := len(children)) else " [empty]"
else:
item_repr = f': {obj}'

if recursive:
recursive_ids.add(objid)

newline_subscription_prefix = f"{prefix}{' ' if last else '| '}"
subscription_repr = _newline_repr(f"{subscription}", newline_subscription_prefix)
yield f"{prefix}{'`- ' if last else '|- '}{subscription_repr}{item_repr}"
prefix += " " if last else "| "
prefix += level_suffix
child_count = len(children)
for index, key, value in children:
yield from _itree(value, subscription=key, prefix=prefix, last=index == (child_count - 1))
53 changes: 42 additions & 11 deletions tests/test_printree.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,21 +7,52 @@
class TestTree(unittest.TestCase):
def test_sorted_with_recursion(self):
"""Confirm multiline key, value and recursion build a similar tree to:
`- . [items=3]
|- A
| B: x
| y
|- B [items=2]
`- . [items=5]
|- A [items=1]
| `- b'xy' [items=1]
| `- (True, False) [items=1]
| `- AB
| CD
| EF
| GH IJ: xx
| y
| y
| zz
|- B [items=1]
| `- 0 [items=1]
| `- b'xy' [items=1]
| `- (True, False) [items=1]
| `- AB
| CD
| EF
| GH IJ: xx
| y
| y
| zz
|- C
| D: x
| y
|- F [items=2]
| |- 0: 1
| `- 1: 2
`- C: <Recursion on dict with id=140712966998864>
`- G [items=2]
|- [items=3]
| |- 0: 1
| |- 1: 2
| `- 2: <Recursion on list with id=1731749379200>
`- .: <Recursion on dict with id=1731749449088>
"""
dct = {"A\nB": "x\ny", "B": (1, "2")}
dct["C"] = dct
multiline = {b"xy": {(True, False): {"AB\nCD\n\tEF\n\t\tGH\tIJ": "xx\ny\ny\n\tzz"}}}
dct = {"A": multiline, "B": (multiline,), "C\nD": "x\ny", "F": (1, "2")}
rec = [1, 2]
rec.append(rec)
dct["G"] = {".": dct, "": rec}
dctid = id(dct)
recid = id(rec)
expected = f"`- . [items=5]\n |- A [items=1]\n | `- b'xy' [items=1]\n | `- (True, False) [items=1]\n | `- AB\n | CD\n | \tEF\n | \t\tGH\tIJ: xx\n | y\n | y\n | \tzz\n |- B [items=1]\n | `- 0 [items=1]\n | `- b'xy' [items=1]\n | `- (True, False) [items=1]\n | `- AB\n | CD\n | \tEF\n | \t\tGH\tIJ: xx\n | y\n | y\n | \tzz\n |- C\n | D: x\n | y\n |- F [items=2]\n | |- 0: 1\n | `- 1: 2\n `- G [items=2]\n |- [items=3]\n | |- 0: 1\n | |- 1: 2\n | `- 2: <Recursion on list with id={recid}>\n `- .: <Recursion on dict with id={dctid}>"
actual = ftree(dct)
expected = '`- . [items=3]\n |- A\n | B: x\n | y\n |- B [items=2]\n | |- 0: 1\n | `- 1: 2\n `- C: <Recursion on dict with id='
self.assertIn(expected, actual)

self.assertEqual(expected, actual)
with patch('sys.stdout', new=StringIO()) as redirected:
ptree(dct) # should be exactly as the ftree result, plus a new line
self.assertEqual(redirected.getvalue(), actual+'\n')

0 comments on commit de3c0b7

Please sign in to comment.