-
Notifications
You must be signed in to change notification settings - Fork 12
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #36 from cidrblock/feature_tree_class
Add tree class
- Loading branch information
Showing
3 changed files
with
321 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,174 @@ | ||
"""An ascii tree generator.""" | ||
from __future__ import annotations | ||
|
||
import os | ||
|
||
from typing import Union | ||
|
||
|
||
class Ansi: | ||
"""ANSI escape codes.""" | ||
|
||
BLUE = "\x1B[34m" | ||
BOLD = "\x1B[1m" | ||
CYAN = "\x1B[36m" | ||
GREEN = "\x1B[32m" | ||
ITALIC = "\x1B[3m" | ||
MAGENTA = "\x1B[35m" | ||
RED = "\x1B[31m" | ||
RESET = "\x1B[0m" | ||
REVERSED = "\x1B[7m" | ||
UNDERLINE = "\x1B[4m" | ||
WHITE = "\x1B[37m" | ||
YELLOW = "\x1B[33m" | ||
|
||
|
||
ScalarVal = Union[bool, str, float, int, None] | ||
JSONVal = Union[ScalarVal, list["JSONVal"], dict[str, "JSONVal"]] | ||
|
||
|
||
class Tree: # pylint: disable=R0902 | ||
"""Renderer for the tree.""" | ||
|
||
PIPE = "│" | ||
ELBOW = "└──" | ||
TEE = "├──" | ||
PIPE_PREFIX = "│ " | ||
SPACE_PREFIX = " " | ||
|
||
def __init__( | ||
self: Tree, | ||
obj: JSONVal, | ||
) -> None: | ||
"""Initialize the renderer.""" | ||
self.obj = obj | ||
self._lines: list[str] = [] | ||
self.blue: list[ScalarVal] = [] | ||
self.bold: list[ScalarVal] = [] | ||
self.cyan: list[ScalarVal] = [] | ||
self.green: list[ScalarVal] = [] | ||
self.italic: list[ScalarVal] = [] | ||
self.magenta: list[ScalarVal] = [] | ||
self.red: list[ScalarVal] = [] | ||
self.reversed: list[ScalarVal] = [] | ||
self.underline: list[ScalarVal] = [] | ||
self.white: list[ScalarVal] = [] | ||
self.yellow: list[ScalarVal] = [] | ||
|
||
def in_color(self: Tree, val: ScalarVal) -> str: | ||
"""Colorize the string. | ||
Args: | ||
val: The thing to colorize | ||
Returns: | ||
The colorized string | ||
""" | ||
if os.environ.get("NO_COLOR"): | ||
return str(val) | ||
ansis = ( | ||
"blue", | ||
"bold", | ||
"cyan", | ||
"green", | ||
"italic", | ||
"magenta", | ||
"red", | ||
"reversed", | ||
"underline", | ||
"white", | ||
"yellow", | ||
) | ||
start = "" | ||
val_str = str(val) | ||
for ansi in ansis: | ||
matches = getattr(self, ansi) | ||
if val_str in [str(match) for match in matches]: | ||
start += getattr(Ansi, ansi.upper()) | ||
return f"{start}{val_str}{Ansi.RESET}" | ||
|
||
@staticmethod | ||
def is_scalar(obj: JSONVal) -> bool: | ||
"""Check if the object is a scalar.""" | ||
return isinstance(obj, (str, int, float, bool)) or obj is None | ||
|
||
def _print_tree( # noqa: C901 | ||
self: Tree, | ||
obj: JSONVal, | ||
prefix: str = "", | ||
is_last: bool = True, # noqa: FBT001, FBT002 | ||
was_list: bool = False, # noqa: FBT001, FBT002 | ||
) -> None: | ||
if isinstance(obj, dict): | ||
if len(obj) > 1: | ||
for key, value in list(obj.items())[:-1]: | ||
_key = f"{Ansi.ITALIC}{key}{Ansi.RESET}" if was_list else key | ||
self.append(f"{prefix}{self.TEE}{self.in_color(_key)}") | ||
self._print_tree( | ||
obj=value, | ||
prefix=prefix + self.PIPE_PREFIX, | ||
is_last=not isinstance(value, (dict, list)), | ||
) | ||
key, value = list(obj.items())[-1] | ||
if was_list: | ||
key = f"{Ansi.ITALIC}{key}{Ansi.RESET}" | ||
self.append(f"{prefix}{self.ELBOW}{self.in_color(key)}") | ||
self._print_tree( | ||
obj=value, | ||
prefix=prefix + self.SPACE_PREFIX, | ||
is_last=True, | ||
) | ||
elif isinstance(obj, list): | ||
if any(isinstance(item, (dict, list)) for item in obj) and len(obj) > 1: | ||
repl_obj = {str(i): item for i, item in enumerate(obj)} | ||
self._print_tree( | ||
obj=repl_obj, | ||
prefix=prefix, | ||
is_last=is_last, | ||
was_list=True, | ||
) | ||
elif isinstance(obj[0], (dict, list)): | ||
self._print_tree(obj=obj[0], prefix=prefix, is_last=True) | ||
elif isinstance(obj[0], (str, int, float, bool)): | ||
for i, item in enumerate(obj): | ||
is_last = i == len(obj) - 1 | ||
_item = str(item) | ||
self.append( | ||
f"{prefix}{self.ELBOW if is_last else self.TEE}{self.in_color(_item)}", | ||
) | ||
else: | ||
err = f"Invalid type in list {type(obj[0])}" | ||
raise TypeError(err) | ||
|
||
elif self.is_scalar(obj): | ||
self.append( | ||
f"{prefix}{self.ELBOW if is_last else self.TEE}{self.in_color(obj)}", | ||
) | ||
else: | ||
err = f"Invalid type {type(obj)}" | ||
raise TypeError(err) | ||
|
||
def append(self: Tree, string: str) -> None: | ||
"""Append a line to the output.""" | ||
self._lines.append(string) | ||
|
||
def render(self: Tree) -> str: | ||
"""Render the root of the tree.""" | ||
if not isinstance(self.obj, dict): | ||
msg = "The root of the tree must be a dict" | ||
raise TypeError(msg) | ||
for k, v in list(self.obj.items())[:-1]: | ||
if isinstance(v, (dict, list)): | ||
self.append(self.in_color(k)) | ||
self._print_tree(v, is_last=not isinstance(v, (dict, list))) | ||
else: | ||
self.append(self.in_color(k)) | ||
self.append(f"{self.ELBOW}{self.in_color(v)}") | ||
k, v = list(self.obj.items())[-1] | ||
if isinstance(v, (dict, list)): | ||
self.append(self.in_color(k)) | ||
self._print_tree(v, is_last=not isinstance(v, (dict, list))) | ||
else: | ||
self.append(self.in_color(k)) | ||
self.append(f"{self.ELBOW}{self.in_color(v)}") | ||
return "\n".join(self._lines) + "\n" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,140 @@ | ||
# cspell:ignore mkey, mfour | ||
"""Test the tree generator.""" | ||
|
||
from __future__ import annotations | ||
|
||
from typing import TYPE_CHECKING | ||
|
||
from pip4a.tree import Tree | ||
|
||
|
||
if TYPE_CHECKING: | ||
import pytest | ||
|
||
from pip4a.tree import JSONVal | ||
|
||
|
||
sample_1: JSONVal = { | ||
"key_one": "one", | ||
"key_two": 42, | ||
"key_three": True, | ||
"key_four": None, | ||
"key_five": ["one", "two", "three"], | ||
"key_six": { | ||
"key_one": "one", | ||
"key_two": 42, | ||
"key_three": True, | ||
"key_four": None, | ||
"key_five": ["one", "two", "three"], | ||
"key_six": { | ||
"key_one": "one", | ||
"key_two": 42, | ||
"key_three": True, | ||
"key_four": None, | ||
"key_five": ["one", "two", "three"], | ||
"key_six": { | ||
"key_one": "one", | ||
"key_two": 42, | ||
"key_three": True, | ||
"key_four": None, | ||
"key_five": ["one", "two", "three"], | ||
}, | ||
}, | ||
}, | ||
} | ||
|
||
result = """key_one | ||
└──one | ||
key_two | ||
└──42 | ||
key_three | ||
└──True | ||
key_four | ||
└──None | ||
key_five | ||
├──one | ||
├──two | ||
└──three | ||
key_six | ||
├──key_one | ||
│ └──one | ||
├──key_two | ||
│ └──42 | ||
├──key_three | ||
│ └──True | ||
├──key_four | ||
│ └──None | ||
├──key_five | ||
│ ├──one | ||
│ ├──two | ||
│ └──three | ||
└──key_six | ||
├──key_one | ||
│ └──one | ||
├──key_two | ||
│ └──42 | ||
├──key_three | ||
│ └──True | ||
├──key_four | ||
│ └──None | ||
├──key_five | ||
│ ├──one | ||
│ ├──two | ||
│ └──three | ||
└──key_six | ||
├──key_one | ||
│ └──one | ||
├──key_two | ||
│ └──42 | ||
├──key_three | ||
│ └──True | ||
├──key_four | ||
│ └──None | ||
└──key_five | ||
├──one | ||
├──two | ||
└──three | ||
""" | ||
|
||
|
||
def test_tree_large(monkeypatch: pytest.MonkeyPatch) -> None: | ||
"""Test the tree generator.""" | ||
monkeypatch.setenv("NO_COLOR", "true") | ||
|
||
assert Tree(sample_1).render() == result | ||
|
||
|
||
sample_2: JSONVal = { | ||
"key_one": True, | ||
"key_two": 42, | ||
"key_three": None, | ||
"key_four": "four", | ||
"key_five": [{"a": 1}, {"b": 2}], | ||
} | ||
|
||
expected = [ | ||
"\x1b[34mkey_one\x1b[0m", | ||
"└──\x1b[32mTrue\x1b[0m", | ||
"\x1b[34mkey_two\x1b[0m", | ||
"└──\x1b[32m42\x1b[0m", | ||
"\x1b[34mkey_three\x1b[0m", | ||
"└──\x1b[32mNone\x1b[0m", | ||
"\x1b[34mkey_four\x1b[0m", | ||
"└──\x1b[32mfour\x1b[0m", | ||
"key_five\x1b[0m", | ||
"├──\x1b[3m0\x1b[0m\x1b[0m", | ||
"│ └──a\x1b[0m", | ||
"│ └──1\x1b[0m", | ||
"└──\x1b[3m1\x1b[0m\x1b[0m", | ||
" └──b\x1b[0m", | ||
" └──2\x1b[0m", | ||
] | ||
|
||
|
||
def test_tree_color() -> None: | ||
"""Test the tree generator.""" | ||
tree = Tree(sample_2) | ||
tree.blue = ["key_one", "key_two", "key_three", "key_four"] | ||
tree.green = [True, 42, None, "four"] | ||
rendered = tree.render().splitlines() | ||
assert rendered == expected |