import os
import pathlib
import sys
from src.fmutils import get_all_dirs
PIPE = "│"
ELBOW = "└──"
TEE = "├──"
PIPE_PREFIX = "│ "
SPACE_PREFIX = " "
[docs]class DirectoryTree:
'''
Parameters
----------
root_dir : string/path
absolute/relative path to root directory containing all files.
dir_only : Bool, optional
whether to only show sub-dirs in the dir-tree. The default is False.
write_tree : Bool, optional
write the full dir-tree in a txt file in current working dir. The default is True.
Returns
-------
None.
'''
def __init__(self, root_dir, dir_only=False, write_tree = True):
self._generator = _TreeGenerator(root_dir, dir_only)
self.write_tree = write_tree
[docs] def generate(self):
tree = self._generator.build_tree()
if self.write_tree:
file = os.getcwd()+'/dir_tree.txt'
with open(file, 'w', encoding="utf-8") as f:
for entry in tree:
f.write("%s\n" % entry)
f.close()
print(f'directory tree file saved at \n {file}')
if not self.write_tree:
for entry in tree:
print(entry)
return
class _TreeGenerator:
def __init__(self, root_dir, dir_only=False):
self._root_dir = pathlib.Path(root_dir)
self._dir_only = dir_only
self._tree = []
def build_tree(self):
self._tree_head()
self._tree_body(self._root_dir)
return self._tree
def _tree_head(self):
self._tree.append(f"{self._root_dir}{os.sep}")
self._tree.append(PIPE)
def _tree_body(self, directory, prefix=""):
entries = self._prepare_entries(directory)
entries_count = len(entries)
for index, entry in enumerate(entries):
connector = ELBOW if index == entries_count - 1 else TEE
if entry.is_dir():
self._add_directory(
entry, index, entries_count, prefix, connector
)
else:
self._add_file(entry, prefix, connector)
def _prepare_entries(self, directory):
entries = directory.iterdir()
if self._dir_only:
entries = [entry for entry in entries if entry.is_dir()]
return entries
entries = sorted(entries, key=lambda entry: entry.is_file())
return entries
def _add_directory(
self, directory, index, entries_count, prefix, connector
):
self._tree.append(f"{prefix}{connector} {directory.name}{os.sep}")
if index != entries_count - 1:
prefix += PIPE_PREFIX
else:
prefix += SPACE_PREFIX
self._tree_body(
directory=directory,
prefix=prefix,
)
self._tree.append(prefix.rstrip())
def _add_file(self, file, prefix, connector):
self._tree.append(f"{prefix}{connector} {file.name}")
[docs]def clone_dir_tree(source_dir, dest_dir):
'''
Parameters
----------
source_dir : string/path
dir form which to clone the dir tree.
dest_dir : string/path
base dir location where the new dir tree will be cloned
Returns
-------
None. Creates the directories at new location without copying files
'''
x = get_all_dirs(source_dir)
dirs = []
[dirs.append(x[j].split('/')[-1]) for j in range(len(x))]
for i in dirs:
try:
os.mkdir(dest_dir + i)
except FileExistsError:
pass
print(f'Cloned {len(x)} directories and sub-directories in total.')
return