Demo: A Directory Tree Generator Tool in Python

In this tutorial, you’ll build a command-line tool tolist the contents of a directory or folder in atreelike diagram. There are already several mature solutions out there that perform this task. You’ll find tools like thetree command, which is available on most operating systems, plus other tools, liketreelib,dirtriex, and so on. However, figuring out your own solution to this problem would be a good learning exercise.

This tutorial refers to the kind of tool described above as adirectory tree generator. The tool you’ll build here will allow you to generate and display a treelike diagram listing the internalstructure of a given directory in yourfile system. You’ll also find this diagram referred to as adirectory tree diagram throughout the tutorial.

Your directory tree generator will have a user-friendly CLI. It’ll also provide some interesting features, such as displaying a tree diagram with the contents of a directory on yourterminal window and saving the diagram to an external file.

Here’s how the application will look and work once you get to the end of this tutorial:

Your directory tree generator will provide a fully functional but minimal CLI with a couple of options that allow you to generate and display a tree diagram listing all the files and directories in a given root directory.

Project Overview

The project you’ll build in this tutorial consists of a command-line application that takes a directory path as an argument, walks through its internal structure, and generates a treelike diagram listing the contents of the directory at hand. In this section, you’ll take a first look at the problem and a possible solution. You’ll also decide how to lay out the project.

./rptree_project/│├── rptree/│   ├── rptree.py│   ├── __init__.py│   └── cli.py│├── README.md└── tree.py

Organizing the Code

In terms ofdesign, if you think of the problem at hand and apply thesingle-responsibility principle, then you can organize the code of your directory tree generator app according to three main responsibilities:

1. Provide the CLI
2. Walk the root directory and build the tree diagram
3. Display the tree diagram

The CLI-related code will live incli.py. Inrptree.py, you’ll place the code related to the second and third responsibilities.

In this example, you’ll write a high-levelDirectoryTree class to generate and display the tree diagram. You’ll use this class in your client code, ormain function. The class will provide a method called.generate() to generate and display the directory tree diagram.

Next, you’ll code a low-level_TreeGenerator class to walk the directory structure and create the list containing the entries that shape the tree diagram. This class will provide a method called.build_tree() to perform this operation.

The tree diagram will have two main components:

1. Head will provide the root directory representation.
2. Body will provide the directory content representation.

The tree head representation will consist of the name of the root directory and an additional pipe (│) character to connect the tree head and body.

The tree body representation will consist ofstrings that include the following components:

• A prefix string that provides the required spacing to reflect the position of an entry in the directory structure
• A character that connects the current subdirectory or file with its parent directory
• The name of the current subdirectory or file

Here’s how you’ll combine these elements to build a directory tree diagram:

Your tree generator’s.build_tree() method will return a list with all the entries that shape the directory tree diagram. To display the diagram, you need to call.generate() on your directory tree object.

Remove ads

Prerequisites

To complete this tutorial and get the most out of it, you should be comfortable with the following concepts:

• Creating command-line interfaces (CLIs) with Python’sargparse module
• Traversing the file system withpathlib
• Using recursion and creating recursive functions in Python
• Working with files usingopen() and thewith statement
• Usingprint() to print text to the screen and also to write to physical files in your file system
• Using object-oriented programming in Python

If you don’t have all of the required knowledge before starting this tutorial, then that’s okay! You can always stop and review the following resources:

• How to Build Command Line Interfaces in Python With argparse
• Python’s pathlib Module: Taming the File System
• Thinking Recursively in Python
• Working With Files in Python
• Your Guide to the Python print() Function
• Object-Oriented Programming (OOP) in Python

In terms of software dependencies, your directory tree generator project doesn’t need any external libraries. All its dependencies are available as Python built-in functions or as modules in the standard library.

That said, it’s time to get your hands dirty with real code and build your own directory tree generator tool!

Step 1: Setting Up the Project Structure

First, you need to create a coherent application layout for your directory tree generator project. Go ahead and create a new directory on your file system with the namerptree_project/. Inside this directory, you need two empty files:

1. README.md
2. tree.py

Next, you need to create a subdirectory calledrptree/ with the following empty files in it:rptree.py,__init__.py, and -cli.py. With this addition, your project’s root directory should look like this:

./rptree_project/│├── rptree/│   ├── rptree.py│   ├── __init__.py│   └── cli.py│├── README.md└── tree.py

To download these files and the code you’ll add to them in this section, click the link below:

At this point, you need an additional setup step. Fire up your favoritecode editor or IDE in your project’s directory, open__init__.py, and add the following content:

# __init__.py"""Top-level package for RP Tree."""__version__="0.1.0"

Python uses__init__.py files to turn a normal directory into a package. Packages contain modules, such asrptree.py andcli.py in this project. Packages and modules are the mechanisms that allow you to organize and structure your Python code.

In this case,__init__.py contains the module’sdocumentation string, commonly known as adocstring. It also defines aglobal constant called__version__, which holds the application’s version number.

Finally, you need a sample directory to test the application and make sure it works correctly. Leave your project’s root directory and create the following directory structure in your file system, side by side with your project’s folder:

../hello/│├── hello/│   ├── __init__.py│   └── hello.py│├── tests/│   └── test_hello.py│├── requirements.txt├── setup.py├── README.md└── LICENSE

This directory structure mimics the general layout of a Python project. You’ll use this sample directory structure to test the directory tree generator tool throughout the steps in this tutorial. This way, you can compare your result with the expected result at any given step on the tutorial.

Step 2: Generating a Directory Tree Diagram in Python

Now that you know the project’s requirements and you’ve set up the project layout and the sample directory, you can start working on the real code. So get your editor ready to jump into coding.

In this section, you’ll code the project’s main functionality. In other words, you’ll write the code to generate a full directory tree diagram from an input directory path. To download that code, click the link below:

Now get back to your code editor and openrptree.py. Then add the following code to the file:

# rptree.py"""This module provides RP Tree main module."""importosimportpathlibPIPE="│"ELBOW="└──"TEE="├──"PIPE_PREFIX="│   "SPACE_PREFIX="    "

In this piece of code, you firstimportos andpathlib from the Python standard library. Next, you define several module-level constants to hold the connector characters and the prefix strings you’ll use to draw the tree diagram on the terminal window. The symbols you’ll use to draw the tree diagram are the same symbols you’ve seen in previous diagrams in this tutorial. The command-line tooltree uses these same symbols to draw tree diagrams.

# rptree.py# Snip...classDirectoryTree:def__init__(self,root_dir):self._generator=_TreeGenerator(root_dir)defgenerate(self):tree=self._generator.build_tree()forentryintree:print(entry)

 1# rptree.py 2# Snip... 3 4class_TreeGenerator: 5def__init__(self,root_dir): 6self._root_dir=pathlib.Path(root_dir) 7self._tree=[] 8 9defbuild_tree(self):10self._tree_head()11self._tree_body(self._root_dir)12returnself._tree1314def_tree_head(self):15self._tree.append(f"{self._root_dir}{os.sep}")16self._tree.append(PIPE)

 1# rptree.py 2# Snip... 3 4class_TreeGenerator: 5# Snip... 6 7def_tree_body(self,directory,prefix=""): 8entries=directory.iterdir() 9entries=sorted(entries,key=lambdaentry:entry.is_file())10entries_count=len(entries)11forindex,entryinenumerate(entries):12connector=ELBOWifindex==entries_count-1elseTEE13ifentry.is_dir():14self._add_directory(15entry,index,entries_count,prefix,connector16)17else:18self._add_file(entry,prefix,connector)

 1# rptree.py 2# Snip... 3 4class_TreeGenerator: 5# Snip... 6 7def_add_directory( 8self,directory,index,entries_count,prefix,connector 9):10self._tree.append(f"{prefix}{connector}{directory.name}{os.sep}")11ifindex!=entries_count-1:12prefix+=PIPE_PREFIX13else:14prefix+=SPACE_PREFIX15self._tree_body(16directory=directory,17prefix=prefix,18)19self._tree.append(prefix.rstrip())2021def_add_file(self,file,prefix,connector):22self._tree.append(f"{prefix}{connector}{file.name}")

>>>fromrptree.rptreeimportDirectoryTree>>>tree=DirectoryTree("../hello")>>>tree.generate()../hello/│├── hello/│   ├── __init__.py│   └── hello.py│├── tests/│   └── test_hello.py│├── requirements.txt├── setup.py├── README.md└── LICENSE

Step 3: Building the Directory Tree Generator’s CLI

There are several tools out there to create CLI applications. Some of the more popular ones areClick,docopt,Typer, and alsoargparse, which is available in the standard library. In your directory tree generator project, you’ll useargparse to provide the command-line interface. This way, you’ll avoid having an external dependency.

Python’sargparse allows you to define the arguments your application will take at the command line and to validate the user’s input. The module also generates help and usage messages for your scripts.

To download the files and the code that you’ll add or modify in this section, click the link below:

To implement the directory tree generator’s CLI, get back to the project’s directory and open thecli.py file from therptree package. Then type in the following code:

"""This module provides the RP Tree CLI."""# cli.pyimportargparseimportpathlibimportsysfrom.import__version__from.rptreeimportDirectoryTreedefmain():args=parse_cmd_line_arguments()root_dir=pathlib.Path(args.root_dir)ifnotroot_dir.is_dir():print("The specified root directory doesn't exist")sys.exit()tree=DirectoryTree(root_dir)tree.generate()

In this piece of code, you first import the required modules from the standard library. Then you import__version__ and alsoDirectoryTree from the containing package,rptree.

Inmain(), you first callparse_cmd_line_arguments() and pack the command line arguments inargs. You’ll see what this function does in a minute. Next, you turn the root directory into apathlib.Path object. The conditional statement does a quick validation to ensure that the user provides a valid directory path and otherwise exits the application.

Finally, you create aDirectoryTree object usingroot_dir as an argument and call.generate() on it to generate and display the corresponding directory tree diagram on your terminal window.

Now you can dive into the code ofparse_cmd_line_arguments(). This function provides all the CLI-related features:

 1# cli.py 2# Snip... 3 4defparse_cmd_line_arguments(): 5parser=argparse.ArgumentParser( 6prog="tree", 7description="RP Tree, a directory tree generator", 8epilog="Thanks for using RP Tree!", 9)10parser.version=f"RP Tree v{__version__}"11parser.add_argument("-v","--version",action="version")12parser.add_argument(13"root_dir",14metavar="ROOT_DIR",15nargs="?",16default=".",17help="Generate a full directory tree starting at ROOT_DIR",18)19returnparser.parse_args()

Here’s what this function does:

• Line 5 instantiatesargparse.ArgumentParser, providing the application’s command name (prog), a shortdescription of the program, and anepilog phrase to display after the user runs the application’s help option. This class provides a parser for all the arguments the user types at the command line.

• Line 10 sets the parser’sversion attribute to a string that holds the application’s name along with its current version,__version__.

• Line 11 adds the firstoptional argument to your application’s CLI. The-v or--version flag is required to provide this argument, which has the default action of displaying the application’s version string on your terminal window.

• Lines 12 to 18 add a second argument to the CLI. Here,root_dir is apositional argument that holds the directory path you’ll use as a starting point to generate the directory tree diagram. In this case, there are four arguments to.add_argument():

  1. metavar holds the name of the argument in usage messages.

  2. nargs defines the number of values your program can take under the argument at hand. For example, your directory tree generator can take only one directory path at the command line, so the appropriate value fornargs is"?".

  3. default provides a default value for the argument at hand. In this case, you use a dot (".") to set the current directory as the default root directory.

  4. help provides a brief help message describing what the argument does.

• Line 19 parses the supplied arguments using.parse_args(). This method returns aNamespace object with all the supplied arguments. You can access these arguments using the dot notation on thenamespace object. Note that you stored this namespace inargs back when you wrotemain().

The final action to complete this step of your journey is to provide an entry-point script. Get back to your code editor and opentree.py, then add the following code to it:

#!/usr/bin/env python3# tree.py"""This module provides RP Tree entry point script."""fromrptree.cliimportmainif__name__=="__main__":main()

This file is short and straightforward. You first importmain() fromcli.py and then wrap its call in the traditionalif __name__ == "__main__": conditional so that Python callsmain() only if you run the file as a program rather than import it as a module.

With this script in place, you can start using your brand-new command-line directory tree generator. Open a command line window, move to the project’s directory, and run the following commands:

$pythontree.py../hello../hello/│├── hello/│   ├── __init__.py│   └── hello.py│├── tests/│   └── test_hello.py│├── requirements.txt├── setup.py├── README.md└── LICENSE$pythontree.py-vRP Tree v0.1.0$pythontree.py--helpusage: tree [-h] [-v] [ROOT_DIR]RP Tree, a directory tree generatorpositional arguments:  ROOT_DIR       Generate a full directory tree starting at ROOT_DIRoptional arguments:  -h, --help     show this help message and exit  -v, --version  show program's version number and exitThanks for using RP Tree!

That’s it! Your directory tree generator tool works. It generates and displays a user-friendly tree diagram on the screen. It also provides version and usage information. That’s pretty cool for about a hundred lines of code! In the next sections, you’ll add a couple more features to the application.

Step 4: Implementing a Directory-Only Option

An interesting feature to add to your directory tree generator is the ability to generate and display directory-only tree diagrams. In other words, a diagram that only displays directories. In this project, you’ll add-d and--dir-only flags to get this done, but before that, you need to update_TreeGenerator so it can support this new feature.

You can download the files and the code that you’ll add or modify in this section by clicking the link below:

Now open therptree.py module and update its code like this:

# rptree.py# Snip...class_TreeGenerator:def__init__(self,root_dir,dir_only=False):self._root_dir=pathlib.Path(root_dir)self._dir_only=dir_onlyself._tree=[]# Snip...def_tree_body(self,directory,prefix=""):entries=self._prepare_entries(directory)entries_count=len(entries)forindex,entryinenumerate(entries):connector=ELBOWifindex==entries_count-1elseTEEifentry.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()ifself._dir_only:entries=[entryforentryinentriesifentry.is_dir()]returnentriesentries=sorted(entries,key=lambdaentry:entry.is_file())returnentries# Snip...

First, you adddir_only as an argument to the class initializer. This is aBoolean argument that allows you to generate a full tree or a directory-only tree depending on the user’s input at the command line. This argument defaults toFalse because generating a full tree is the most common use case.

In the second highlighted line, you create an instance attribute called._dir_only to hold the newly added argument.

In the third highlighted line, you replace two lines of the original code with a call to._prepare_entries(). As its name suggests, this function prepares the directory entries to generate either a full tree or a directory-only tree.

In._prepare_entries(), you first get theentriesgenerator. Theif statement checks if._dir_only isTrue. If so, then you filter out the files with alist comprehension and return alist of directories. If._dir_only isFalse, then you sort the entries, reusing the same code you saw before. Finally, you return the complete list of entries indirectory.

Now you need to make sure that you pass this new argument to the instance of_TreeGenerator back inDirectoryTree:

# rptree.py# Snip...classDirectoryTree:def__init__(self,root_dir,dir_only=False):self._generator=_TreeGenerator(root_dir,dir_only)# Snip...

In the first highlighted line, you add a new argument calleddir_only to the class initializer. In the second highlighted line, you make sure to pass the new argument to the constructor of_TreeGenerator.

With these changes in place, you can update thecli.py file so that the application can take and process the-d and--dir-only flags at the command line. First, you need to updatemain():

# cli.py# Snip...defmain():# Snip...tree=DirectoryTree(root_dir,dir_only=args.dir_only)tree.generate()

In the highlighted line, you passargs.dir_only to thedir_only argument ofDirectoryTree. This attribute of theargs namespace holds a Boolean value that depends on the user’s input. If the user provides the-d or--dir-only option at the command line, thenargs.dir_only isTrue. Otherwise, it’sFalse.

Next, go and add those-d and--dir-only flags to the command-line interface. To do that, you need to updateparse_cmd_line_arguments() like this:

# cli.py# Snip...defparse_cmd_line_arguments():# Snip...parser.add_argument("-d","--dir-only",action="store_true",help="Generate a directory-only tree",)returnparser.parse_args()

Theaction argument in the call to.add_argument() holds the value"store_true", which means that this argument automatically storesTrue orFalse according to the user’s input. In this case, if the user provides the-d or--dir-only flag at the command line, then the argument storesTrue. Otherwise, it storesFalse.

With this update in place, it’s time to run and test the application. Get back to your terminal window and execute the following command:

$pythontree.py../hello-d../hello/│├── hello/│└── tests/

From this point on, if you provide the-d or-dir-only flag at the command line, then the tree diagram only displays the subdirectories in your samplehello/ directory.

Step 5: Saving the Directory Tree Diagram to a File

In this section, you’ll add a final feature to your directory tree generator tool. You’ll provide the app with the capability to save the generated directory tree diagram to an external file. To do that, you’ll add a new argument to the CLI with the flags-o and--output-file.

As usual, to download the code that you’ll add or modify in this section, click the link below:

Now go back torptree.py and updateDirectoryTree like this:

# rptree.py# Snip...importsys# Snip...classDirectoryTree:def__init__(self,root_dir,dir_only=False,output_file=sys.stdout):self._output_file=output_fileself._generator=_TreeGenerator(root_dir,dir_only)defgenerate(self):tree=self._generator.build_tree()ifself._output_file!=sys.stdout:# Wrap the tree in a markdown code blocktree.insert(0,"```")tree.append("```")self._output_file=open(self._output_file,mode="w",encoding="UTF-8")withself._output_fileasstream:forentryintree:print(entry,file=stream)

This update is almost a full reimplementation ofDirectoryTree. First, you add a new argument to the class initializer calledoutput_file. This argument defaults tosys.stdout, which is the standard output (your screen). Then you store the newly added argument in an instance attribute called._output_file.

In.generate(), you first build the directory tree diagram and store it intree. The conditional statement checks if the user has provided an output file different fromsys.stdout. If so, then theif code block wraps the tree diagram in a markdown code block using backticks ("```").

Next, you open the provided output file usingopen() so you can process it using thewith statement.

In thewith block, you start afor loop to print the directory tree diagram to the provided output file. Note thatprint() can also write to regular files on your file system. To do this, you just need to provide a customfile argument. To dive deeper into the features ofprint(), check outYour Guide to the Python print() Function.

Once you’ve finished withDirectoryTree, you can update the command-line interface to enable the output file option. Get back tocli.py and modify it like this:

# cli.py# Snip...defmain():# Snip...tree=DirectoryTree(root_dir,dir_only=args.dir_only,output_file=args.output_file)tree.generate()defparse_cmd_line_arguments():# Snip...parser.add_argument("-o","--output-file",metavar="OUTPUT_FILE",nargs="?",default=sys.stdout,help="Generate a full directory tree and save it to a file",)returnparser.parse_args()

The first step is to take the output file as an argument in theDirectoryTree constructor. The output file, if any, will be stored inargs.output_file.

Next, you add a new argument toparser. This argument has two flags:-o and--output-file. To provide an alternative output file, the user has to use one of these flags and provide the path to the files at the command line. Note that the output file defaults tosys.stdout. This way, if the user doesn’t provide an output file, then the application automatically uses the standard output, the screen.

You can test the newly added option by running the following command on your terminal:

$pythontree.py../hello-ooutput_file.md

This command generates a full directory tree diagram and saves it into theoutput_file.md file in your current directory. If you open the file, then you’ll see the directory tree diagram saved there in markdown format.

That’s it! Your directory tree generator project is complete. Besides the default option that generates and displays a full directory tree diagram, the application provides the following options:

• -v,--version show the current version information and exit the application.
• -h,--help show help and usage messages.
• -d,--dir-only generate a directory-only tree and print it into the screen.
• -o,--output-to-markdown generate a tree and save it to a file in markdown format.

You now have a fully functional command-line tool that generates user-friendly directory tree diagrams. Great job!

Conclusion

You can automate and speed up several processes and tasks in your working environment by creatingCLI tools and applications. In Python, you can quickly create this kind of tool usingargparse or other third-party libraries. In this tutorial, you wrote a full project to build a Pythondirectory tree generator tool for your command line.

The application takes a directory path at the command line, generates a directory tree diagram, and displays it on your terminal window or saves it to an external file on your file system. It also provides a few more options to tweak the resulting tree diagram.

In this tutorial, you learned how to:

• Create aCLI application with Python’sargparse
• Recursivelytraverse a directory structure usingpathlib
• Generate, format, and print adirectory tree diagram
• Save the directory tree diagram to anoutput file

The final source code for the directory tree generator project is available for you to download. To get it, click the link below:

Next Steps

Up to this point, you’ve built a fully functional directory tree generator tool. Even though the application provides a minimal set of features, it’s a good starting point for you to continue adding features and learning in the process. This will help you take your skills with Python and CLI applications to the next level.

Here are a few ideas you can implement to continue improving your directory tree generator tool:

• Add support for sorting files and directories: The ability to sort files and directories is a great feature to have. For example, you can add-s and--sort-tree Boolean flags to allow the user to tweak the order of files and directories in the final tree diagram.

• Add icons and colors to the tree diagram: Adding icons, font colors, or both is also a nice feature to implement. For example, you can use custom folder icons for the directories and file type–based icons for the files.

• Set up the application to publish it as an open source project: Preparing the application to publish to PyPI as an open source project might be an interesting challenge for you to take. Doing so will allow you to share your work with your friends and colleagues. To get started with publishing packages to PyPI, check outHow to Publish an Open-Source Python Package to PyPI.

These are just a few ideas of how you can continue adding features to your directory tree generator. Take the challenge and build something amazing on top of this!