How to generate a Python directory list (tree structure)?

Question

I know that we can use os.walk() to list all sub-directories or all files in a directory. How can I generate a tree structure as follows rather than just listing all the directories? In some sub-directories I have hundreds of files, so how can I show only few files within the sub-directory without listing them all?

doc/
├── _static/
│   ├── embedded/
│   │   ├── deep_file
│   │   └── very/
│   │       └── deep/
│   │           └── folder/
│   │               └── very_deep_file1
|   |               └── very_deep_file2
|   |               └── ......
│   └── less_deep_file
├── about.rst
├── conf.py
└── index.rst

Answer 1

Stumbled upon this and decided to try it for fun, I hope it helps someone. Tested on Windows 10, python 3.9.10

import os


class FileSystemItem(object):
  def __init__(self, path: str, parent: 'FileSystemItem' = None):
    """Base class
    @param path: path to this Item, can be a [File, Folder]
    @param parent: parent of this item, meant to be a Folder. If None, this Item is a root item"""

    super(FileSystemItem, self).__init__()

    self.parent = parent
    self.path = os.path.abspath(path)
    self.name = os.path.split(self.path)[1]

    self.flag_isroot = not isinstance(parent, FileSystemItem)
    self.flag_isdir = False
    self.flag_isfile = False
    self.flag_islast = False

  def set_name(self, name: str):
    self.name = name

  def get_parent(self):
    return self.parent

  def get_path(self):
    return self.path

  def get_name(self):
    return self.name

  def set_islast(self, flag: bool):
    self.flag_islast = flag

    return self

  def islast(self):
    return self.flag_islast

  def isroot(self):
    return self.flag_isroot

  #/////////////////////NOTE NOTE NOTE NOTE NOTE NOTE////////////////////////////
  #///////////////// YOU MAY WANT TO WORK ON THE SORTING...//////////////////////
  #////if You don't want things like Name 1, Name 10, Name 2, Name 20, Name 3////
  #//And you want things to look like Name 1, Name 2, Name 3, Name 10, Name 20//
  #//////////////////////////////////////////////////////////////////////////////

  def __eq__(self, other: 'FileSystemItem'):
    if all([self.flag_isdir, other.flag_isdir]) or all([self.flag_isfile, other.flag_isfile]):
      return self.name == other.name

    return False

  def __ne__(self, other: 'FileSystemItem'):
    if all([self.flag_isdir, other.flag_isdir]) or all([self.flag_isfile, other.flag_isfile]):
      return self.name != other.name

    return True

  def __lt__(self, other: 'FileSystemItem'):
    if all([self.flag_isdir, other.flag_isdir]) or all([self.flag_isfile, other.flag_isfile]):
      return self.name < other.name

    return self.flag_isdir

  def __gt__(self, other: 'FileSystemItem'):
    if all([self.flag_isdir, other.flag_isdir]) or all([self.flag_isfile, other.flag_isfile]):
      return self.name > other.name

    return self.flag_isfile

  def __le__(self, other: 'FileSystemItem'):
    if all([self.flag_isdir, other.flag_isdir]) or all([self.flag_isfile, other.flag_isfile]):
      return self.name <= other.name

    return self.flag_isdir

  def __ge__(self, other: 'FileSystemItem'):
    if all([self.flag_isdir, other.flag_isdir]) or all([self.flag_isfile, other.flag_isfile]):
      return self.name >= other.name

    return self.flag_isfile

  def __str__(self):
    return self.path


class File(FileSystemItem):
  def __init__(self, path: str, parent: 'FileSystemItem' = None):
    """Create file item to represent a file in storage
    @param path: path to the file in question
    @param parent: parent of this file, meant to be a Folder."""

    super(File, self).__init__(path, parent)

    self.flag_isfile = True


class Folder(FileSystemItem):
  def __init__(self, path: str, parent: 'FileSystemItem' = None, filters: tuple = ('.*', )):
    """Create folder item to represent a folder in storage
    @param path: path to the folder in question
    @param parent: parent of this folder, meant to be a Folder. If None, this Folder is a root folder
    @param filters: tuple containing file extensions. The folder item will only contain files that have these extensions"""

    super(Folder, self).__init__(path, parent)

    self.flag_isdir = True
    # self.flag_isroot = parent is None

    self.contents = {'folders': [], 'files': []}

    if not filters: filters = ('.*',)

    for path in [f'{self.path}\\{item}' for item in os.listdir(self.path)]:
      try:
        if os.path.isdir(path):
          self.contents['folders'].append(Folder(path, self, filters))
        else:
          if '.*' in filters:
            self.contents['files'].append(File(path, self))
          elif path.lower().endswith(filters):
            self.contents['files'].append(File(path, self))
      except PermissionError:
        pass

  def get_contents(self, files_limit: int = None):
    """get this folder's items [Folders and/or not Files]
    @param files_limit: number of Files items to fetch. If a lower number than Files count is provided, a dummy File is appended at the end of the list"""

    items = [FileSystemItem.set_islast(item, False) for item in self.contents['folders']]
    if files_limit is None:
      items.extend([FileSystemItem.set_islast(item, False) for item in self.contents['files']])
      items.sort()
    else:
      if (files_l := self.contents['files']):
        limit = files_limit if 0 <= files_limit <= len(files_l) else len(files_l)
        items.extend([FileSystemItem.set_islast(files_l[i], False) for i in range(limit)])
        items.sort()
        if not limit == len(files_l):
          rem = len(files_l) - limit
          dummy_file_obj = File('*root:/some/dummy/path/to/some/dummy/file', self)
          dummy_file_obj.set_name(' ......')
          # dummy_file_obj.set_name(' \033[92m+{rem} more\033[0m....'.format(rem=rem))
          items.append(dummy_file_obj)

    if items: items[-1].set_islast(True)

    return items

  def get_folders(self):
    """get Folders items (only) of this Folder"""

    items = [FileSystemItem.set_islast(item, False) for item in self.contents['folders']]
    items.sort()

    if items: items[-1].set_islast(True)

    return items

  def get_files(self):
    """get Files items (only) of this Folder"""

    items = [FileSystemItem.set_islast(item, False) for item in self.contents['files']]
    items.sort()

    if items: items[-1].set_islast(True)

    return items

  def tree_structure(self, include_files: bool = True, files_limit: int = None):
    """get tree structure of this Folder
    @param include_files: decide whether to show files in the tree
    @param files_count: if show files, decide how many files to show per branch"""

    def get_prefix(item: FileSystemItem):
      prefix = '' if item.isroot() else '├───' if not item.islast() else '└───'
      
      while not (item := item.get_parent()) is None and not item.isroot():
        prefix = f"{' ' * 4}{prefix}" if item.islast() else f"│{' ' * 3}{prefix}"
      
      return prefix

    def get_tree(start: Folder, output: list[str] = []):
      output.append(f"{get_prefix(start)}{start.get_name()}\n")

      if isinstance(start, Folder):
        for item in start.get_contents(files_limit) if include_files else start.get_folders():
          get_tree(item, output)
      
      return output

    return "".join(get_tree(self))


if __name__ == '__main__':
  os.system('color')

  root = Folder(path='D:/VIDEOS', parent=None, filters=('.*',))

  print(root.tree_structure(include_files=True, files_limit=3))

  with open((save_at := os.path.abspath('tree_structure.txt')), mode='w', encoding='utf-8') as file:
    file.seek(0); file.truncate()
    file.write(root.tree_structure(include_files=True, files_limit=None))

    print(f"[INFO ] [tree structure saved at '{save_at}' ]")

Answer 2

The following may be help you:

for directory, _, filenames in os.walk('your path'):
    for filename in filenames:
        print(os.path.join(directory, filename))

Answer 3

os.walk gives you a root, directories, files when you iterate over it. Why not simply loop over files[:5] and then print an extra ... if the length exceeds 5?