A reimplementation of the `tree` command using rich.tree

pytree.py

#!/usr/bin/env python3
import sys
import argparse
from pathlib import Path
from typing import Optional

try:
    from rich.console import Console
    from rich.tree import Tree
except ImportError:
    # if rich is not installed try using rich bundled with pip
    try:
        from pip._vendor.rich.console import Console
        from pip._vendor.rich.tree import Tree
    except ImportError:
        print('Error: Failed to load dependencies !!!', file=sys.stderr)
        print('You are using a version of python that is inferion than 3.7 or a version of pip inferion to 22.0', file=sys.stderr)
        print('Try upgrading pip or your python version, or install the rich library', file=sys.stderr)
        sys.exit(1)


error_console = Console(stderr=True, style="bold red")
std_console = Console()


def parse_args(argv=None):
    """ Parsing arguments """
    argv = argv if argv is not None else sys.argv[1:]
    parser = argparse.ArgumentParser(description="List files & directories in a tree like structure.")
    parser.add_argument('dirname')
    parser.add_argument('-L', '--max-depth', default=None, help="maximum depth")
    parser.add_argument('-p', '--pattern', default=None, help="maximum depth")

    return parser.parse_args(argv)


def add_childen(
    location: Path, root_node, depth: int = 0, max_depth: Optional[int] = None,
    file_pattern: Optional[str] = None
):
    """ Add childer to node """
    items = [i for i in location.iterdir()]
    folders = [i for i in items if i.is_dir()]

    if max_depth is not None and depth >= max_depth:
        return

    for cur_dir in folders:
        folder_node = root_node.add(f":file_folder: {cur_dir.name}/")
        add_childen(cur_dir, folder_node, depth=depth + 1, file_pattern=file_pattern, max_depth=max_depth)

    if file_pattern is not None:
        files = [i for i in location.glob(file_pattern)]
    else:
        files = [i for i in items if i.is_file()]

    for f in files:
        root_node.add(f"📄 {f.name}")


def tree_cmd(argv=None):
    args = parse_args(argv)

    root_dir = Path(args.dirname)
    max_depth = args.max_depth
    if max_depth:
        max_depth = int(max_depth)

    if not root_dir.is_dir():
        error_console.print(f'{root_dir} is not a directory or does not exist !!!')
        sys.exit(1)
    
    root_node = Tree(f":file_folder: {root_dir.resolve().name}", highlight=True)
    add_childen(root_dir, root_node, max_depth=max_depth, file_pattern=args.pattern)

    std_console.print(root_node)


if __name__ == "__main__":
    tree_cmd()