FileTree.py

#!python
import os
import pathlib
from rich import print
import mimetypes

mimetypes.init()
PIPE = "│"
ELBOW = "└──"
TEE = "├──"
PIPE_PREFIX = "│ "
SPACE_PREFIX = " "


class TreeGenerator:
 ext = None

 def __init__(self, root_dir, file_extension):
 self.root_dir = pathlib.Path(root_dir)
 self.tree = []
 if file_extension:
 self.ext = tuple(
 [
 exten if exten.startswith(".") else "." + exten
 for exten in file_extension
 ]
 )

 # Builder function used to build the entire directory tree
 # Main thing to call
 def build_tree(self):
 self.tree_head()
 self.tree_body(self.root_dir)
 return self.tree

 # Binds the main directory in which the program is called
 # This goes at the top of the directory
 def tree_head(self):
 self.tree.append(f"{self.root_dir}{os.sep}")
 self.tree.append(PIPE)

 # After the root each directories and their traversal are appended here
 # The iterdir function iteratest the present directory giving the list of all items
 # in it
 def tree_body(self, directory, prefix=""):
 entries = directory.iterdir()
 # The items so obtained are sorted depending on filetype
 # i.e. if the directory has files then they should be listed
 # prior to the directories
 # The files so obtained are connected via connector
 # If its last file end it with elbow else continue with a tee
 entries = sorted(entries, key=lambda fentry: fentry.is_file())
 entries_count = len(entries)
 for index, entry in enumerate(entries):
 try:
 connector = ELBOW if index == entries_count - 1 else TEE
 # If the type of file is directory then append it using directory type
 # connection whilst traversing it again ahead recursively
 if entry.is_dir():
 self.add_directory(entry, index, entries_count, prefix, connector)
 # Else the type is file so append it using file type connection throughout
 else:
 self.add_file(entry, prefix, connector)
 except PermissionError:
 continue

 def add_directory(self, directory, index, entries_count, prefix, connector):
 # First the name of the directory is appended before traversing
 self.tree.append(
 f"{prefix}{connector} [bright_yellow]{directory.name}[/]{os.sep}"
 )
 # if its not last element (file element) then provide a gap in the nesting so formed
 # ex if we have dir1, dir2 in fold which is in pfold
 # pfold
 # │ └──fold
 # │<----->├── dir1
 # │<----->└── dir2
 # Else if its last element then just escape it out with a gap
 # ex if we have dir1, dir2 in fold which is in pfold and then next fold in pfold is fold2
 # pfold
 # │ └──fold
 # │ ├── dir1
 # │ └── dir2
 # │<-----> last element gap
 # ├── fold2

 if index != entries_count - 1:
 prefix += PIPE_PREFIX
 else:
 prefix += SPACE_PREFIX
 # Recursively traverse everything down the directory
 self.tree_body(
 directory=directory,
 prefix=prefix,
 )
 self.tree.append(prefix.rstrip())

 def add_file(self, file, prefix, connector):
 # Obtain file types just for pretty coloring out different file
 # types in different colors
 ftype = mimetypes.guess_type(file)[0]
 # Check if user specified any file extensions
 # helps to locate out specific file types thoughout
 # your search directory
 if self.ext:
 if file.name.endswith(self.ext):
 if ftype is not None:
 match ftype.split("/")[0]:
 # If file type is media highlight it with blue
 case "image" | "audio" | "video":
 self.tree.append(
 f"{prefix}{connector} [blue]{file.name}[/]"
 )
 # If file type is executable highlight it with red
 case "application":
 self.tree.append(f"{prefix}{connector} [red]{file.name}[/]")
 # If file type is source or text highlight it with green
 case "text":
 self.tree.append(
 f"{prefix}{connector} [spring_green3]{file.name}[/]"
 )
 # If file type is unknown to os then grey it out
 else:
 self.tree.append(
 f"{prefix}{connector} [bright_black]{file.name}[/]"
 )
 else:
 if ftype is not None:
 match ftype.split("/")[0]:
 # If file type is media highlight it with blue
 case "image" | "audio" | "video":
 self.tree.append(f"{prefix}{connector} [blue]{file.name}[/]")
 # If file type is executable highlight it with red
 case "application":
 self.tree.append(f"{prefix}{connector} [red]{file.name}[/]")
 # If file type is source or text highlight it with green
 case "text":
 self.tree.append(
 f"{prefix}{connector} [spring_green3]{file.name}[/]"
 )
 # If file type is unknown to os then grey it out
 else:
 self.tree.append(f"{prefix}{connector} [bright_black]{file.name}[/]")


# Another class more abstraction ¯\(°_o)/¯
class DirectoryTree:
 def __init__(self, root_dir, *extension):
 self._generator = TreeGenerator(root_dir, extension)

 # This generates the tree
 def generate(self):
 tree: list[str] = self._generator.build_tree()
 write = input("Want to write it to file (y/n):").lower() == "y"
 fname = None
 if write:
 fname = input(
 "Enter full path to file (Leaving empty will create a temp.txt file where this prog executes)\n File Name: "
 )
 if fname is not None and os.path.exists(fname) and os.path.isfile(fname):
 fpoint = open(fname, "w")
 else:
 fpoint = open("temp.txt", "w+")
 for entry in tree:
 print(entry, file=fpoint)
 print(f"└───[bright_white] Done ────[/]", file=fpoint)
 fpoint.close()
 else:
 for entry in tree:
 print(entry)
 print(f"└───[bright_white] Done ────[/]")


dir_path = input("Enter absolute path to directory: ")
ext = input(
 "Enter file extension(s) separated with space. Leave blank to list everything: "
)
if ext != "":
 ext = ext.split(" ")
else:
 ext = None
if os.path.exists(dir_path) and os.path.isdir(dir_path):

 if ext is not None:
 f_tree = DirectoryTree(dir_path, *ext)
 f_tree.generate()
 else:
 f_tree = DirectoryTree(dir_path)
 f_tree.generate()
else:
 print("Entered directory doesn't exists")