diff --git a/.vscode/launch.json b/.vscode/launch.json index 6ace40d..aac62a6 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -11,6 +11,13 @@ "module": "pip4a", "args": ["inspect", "-v", "-v", "-vv"], "justMyCode": true + }, + { + "name": "Python: Current File", + "type": "python", + "request": "launch", + "program": "${file}", + "console": "integratedTerminal" } ] } diff --git a/src/pip4a/tree.py b/src/pip4a/tree.py new file mode 100644 index 0000000..5287353 --- /dev/null +++ b/src/pip4a/tree.py @@ -0,0 +1,177 @@ +"""An ascii tree generator.""" +from __future__ import annotations + +import os + + +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 = bool | str | float | int | None +JSONVal = ScalarVal | list["JSONVal"] | dict[str, "JSONVal"] + + +class Tree: # pylint: disable=R0902 + """Renderer for the tree.""" + + PIPE = "│" + ELBOW = "└──" + UEBLOW = "┌──" + 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: + string: The string to colorize + """ + if os.environ.get("NO_COLOR"): + return str(val) + ansis = ( + "blue", + "bold", + "cyan", + "green", + "italic", + "magenta", + "red", + "reversed", + "underline", + "white", + "yellow", + ) + start = "" + if val == "four": + pass + for ansi in ansis: + matches = getattr(self, ansi) + try: + index = matches.index(val) + except ValueError: + continue + + if isinstance(val, type(matches[index])): + start += getattr(Ansi, ansi.upper()) + + return f"{start}{val}{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" diff --git a/tests/unit/test_tree.py b/tests/unit/test_tree.py new file mode 100644 index 0000000..d249e9c --- /dev/null +++ b/tests/unit/test_tree.py @@ -0,0 +1,138 @@ +"""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: 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).render() == result + + +sample = { + "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) + tree.blue = ["key_one", "key_two", "key_three", "key_four"] + tree.green = [True, 42, None, "four"] + rendered = tree.render().splitlines() + assert rendered == expected