flatten output of tree(1)

untree.py

#!/usr/bin/python3
# encoding: utf-8
# SPDX-FileCopyrightText: 2023 FC Stegerman <flx@obfusk.net>
# SPDX-License-Identifier: GPL-3.0-or-later

from typing import Iterable, Iterator, List, Tuple

INDENT_SPACE = {"    ", "│   ", "|   "}
INDENT_FINAL = {"└── ", "├── ", "`-- ", "|-- "}


def untree(lines: Iterable[str]) -> Iterator[Tuple[str, ...]]:
    r"""
    Flatten output of tree(1).

    NB: stops processing when it encounters a blank line.

    >>> tree = '''
    ... a
    ... └── b
    ...     ├── c
    ...     │   ├── bar
    ...     │   └── d
    ...     │       └── baz
    ...     └── foo
    ...
    ... 4 directories, 3 files
    ... '''[1:-1]
    >>> for path in untree(tree.splitlines()):
    ...     print("/".join(path))
    a
    a/b
    a/b/c
    a/b/c/bar
    a/b/c/d
    a/b/c/d/baz
    a/b/foo

    """
    path: List[str] = []
    for line in lines:
        if not (line := line.rstrip("\r\n")):
            break
        indent, name = _split_indent(line)
        path = path[:indent] + [name]
        yield tuple(path)


def _split_indent(s: str) -> Tuple[int, str]:
    i, t = 0, s
    while any(t.startswith(x) for x in INDENT_SPACE):
        i, t = i + 1, t[4:]
    if any(t.startswith(x) for x in INDENT_FINAL):
        return i + 1, t[4:]
    return 0, s


if __name__ == "__main__":
    import sys
    for path in untree(sys.stdin):
        print("/".join(path))

# vim: set tw=80 sw=4 sts=4 et fdm=marker :