show_dir

#!/usr/bin/python3

import os
import sys
from subprocess import check_output, Popen, PIPE
from termcolor import colored
from optparse import OptionParser

FLAG = ""

def format_dirs(dirs:str) -> list:
    """
        Turns raw data from the 'ls' command
        into a list of files and directories.
    """
    dirs = dirs.split("\\n")
    dirs[0] = dirs[0][2:]
    dirs = dirs[:-1]

    return dirs


def get_local(path:str, indent:int) -> None:
    """
        Discovers and prints the file tree.
    """
    local_dirs = []
    try:
        if os.path.isdir(path):
            os.chdir(path)
            if FLAG:
                dirs = format_dirs(str(check_output(["ls", FLAG, path])))
            else:
                dirs = format_dirs(str(check_output(["ls", path])))
        else:
            return
    except Exception as e:
        sys.exit("Exception during get_local" + e)

    for dir in dirs:
        os.chdir(path)
        if os.path.isfile(dir):
            print_file(dir, indent)
        elif os.path.isdir(dir):
            local_dirs.append(dir)

    for dir in local_dirs:
        print_folder(dir, indent)
        get_local(path + "/" + dir, indent+1)


def print_file(file:str, indent) -> None:
    """
        Prints file to the screen in tree
        structured formatting.
    """
    if indent > 0:
        print("|    "*(indent-1) + "|    |---> " + colored(file, "green"), file = sys.stdout)
    else:
        print("|    "*(indent-1) + "|---> " + colored(file, "green"), file = sys.stdout)


def print_folder(folder:str, indent) -> None:
    """
        Prints folder to the screen in tree
        structured formatting.
    """
    print("|    "*indent + "|", file = sys.stdout)
    print("|    "*indent + "|\\____" + colored(folder, "cyan"), file = sys.stdout)


def setup_parser() -> OptionParser:
    """
        Creates, sets-up and returns an
        OptionParser object.
    """
    parser = OptionParser(usage="%prog [-s <file>] <start dir>",
                          version="1.2.0",
                          description= """Displays the files and folders located under <start dir>,
                                          if no start dir is specified the working directory will be used.
                                       """
                         )

    parser.add_option("-a",
                      action="store_true",
                      dest="ALL_FILES",
                      default=False,
                      help="include hidden files"
                     )

    user = os.environ.get("USER")
    if user.lower() == "root":
        user = os.environ.get("SUDO_USER")
    parser.add_option("-s",
                      metavar="<file>",
                      dest="save_file",
                      default=None,
                      help  = f"save the tree to /home/{user}/Documents/<file>."
                     )

    return parser


def main() -> None:

    parser = setup_parser()

    (opts, args) = parser.parse_args()

    if len(args) > 1:
        parser.print_help()
        sys.exit(1)
    elif len(args) == 1:
        start_dir = args[0]
    else:
        start_dir = os.getcwd()

    if start_dir.endswith("/") and len(start_dir) > 1:
        start_dir = start_dir[:-1]

    if not os.path.isdir(start_dir):
        print(colored(f"{start_dir} is not a correct path", "red"))
        parser.print_help()
        sys.exit(1)

    global FLAG
    if opts.ALL_FILES:
        FLAG = "-A"

    global stdout
    if opts.save_file is None:
        sys.stdout = sys.__stdout__
    else:
        user = os.environ.get("USER")
        if user.lower() == "root":
            user = os.environ.get("SUDO_USER")
        sys.stdout = open(f"/home/{user}/Documents/{opts.save_file}","w")

    print(colored(start_dir + "/", "cyan"), file = sys.stdout)
    get_local(start_dir, 0)


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f"Error: {e}.", file = sys.stderr)
    finally:
        if sys.stdout != sys.__stdout__:
            sys.stdout.close()
        sys.stdout = sys.__stdout__