Unix `tree` command output parser

parse_tree_output.py

from dataclasses import dataclass
import sys
from typing import List, TextIO


@dataclass
class TreeNode:
    name: str
    children: List["TreeNode"]

    def __repr__(self) -> str:
        return f"TreeNode(name={self.name!r}, children=[{len(self.children)} nodes])"

    def __getitem__(self, key: str) -> "TreeNode":
        try:
            return next(child for child in self.children if child.name == key)
        except StopIteration:
            raise KeyError(key)


def parse_tree_output(io: TextIO) -> List[TreeNode]:
    """ Parse output of Unix `tree` command. """
    result: List[TreeNode] = []
    stack: List[TreeNode] = []
    for line in io:
        line = line.rstrip()
        if not line:
            break
        if "├──" not in line and "└──" not in line:
            node = TreeNode(name=line, children=[])
            result.append(node)
            stack = [node]
            continue
        level = line.count("│")
        node = TreeNode(name=line[level*4+4:], children=[])
        stack[level].children.append(node)
        stack[level+1:] = [node]
    return result


if __name__ == "__main__":
    print(parse_tree_output(sys.stdin))