Travel directory tree with limited recursion depth

Question

I need to process all files in a directory tree recursively, but with a limited depth.

That means for example to look for files in the current directory and the first two subdirectory levels, but not any further. In that case, I must process e.g. ./subdir1/subdir2/file, but not ./subdir1/subdir2/subdir3/file.

How would I do this best in Python 3?

Currently I use os.walk to process all files up to infinite depth in a loop like this:

for root, dirnames, filenames in os.walk(args.directory):
    for filename in filenames:
        path = os.path.join(root, filename)
        # do something with that file...

I could think of a way counting the directory separators (/) in root to determine the current file's hierarchical level and break the loop if that level exceeds the desired maximum.

I consider this approach as maybe insecure and probably pretty inefficient when there's a large number of subdirectories to ignore. What would be the optimal approach here?

Answer 1

Starting in python 3.5, os.scandir is used in os.walk instead of os.listdir. It works many times faster. I corrected @kevin sample a little.

import os

def walk(top, maxdepth):
    dirs, nondirs = [], []
    for entry in os.scandir(top):
        (dirs if entry.is_dir() else nondirs).append(entry.path)
    yield top, dirs, nondirs
    if maxdepth > 1:
        for path in dirs:
            for x in walk(path, maxdepth-1):
                yield x

for x in walk(".", 2):
    print(x)