List directories with a specified depth in Python

Question

I\'m want a function to return a list with directories with a specified path and a fixed depth and soon realized there a few alternatives. I\'m using os.walk quite a lot but

Answer 1

A simple, recursive solution using os.scandir:

def _walk(path, depth):
    """Recursively list files and directories up to a certain depth"""
    depth -= 1
    with os.scandir(path) as p:
        for entry in p:
            yield entry.path
            if entry.is_dir() and depth > 0:
                yield from _walk(entry.path, depth)

Answer 2

If the depth is fixed, glob is a good idea:

import glob,os.path
filesDepth3 = glob.glob('*/*/*')
dirsDepth3 = filter(lambda f: os.path.isdir(f), filesDepth3)

Otherwise, it shouldn't be too hard to use os.walk:

import os,string
path = '.'
path = os.path.normpath(path)
res = []
for root,dirs,files in os.walk(path, topdown=True):
    depth = root[len(path) + len(os.path.sep):].count(os.path.sep)
    if depth == 2:
        # We're currently two directories in, so all subdirs have depth 3
        res += [os.path.join(root, d) for d in dirs]
        dirs[:] = [] # Don't recurse any deeper
print(res)

Answer 3

This is not exactly neat, but under UNIX-like OS, you could also rely on a system tool like "find", and just execute it as an external program, for example:

from subprocess import call
call(["find", "-maxdepth", "2", "-type", "d"])

You can then redirect the output to some string variable for further handling.

Answer 4

I really like phihag's answer. I adapted it to suit my needs.

import fnmatch,glob
def fileNamesRetrieve( top, maxDepth, fnMask  ):
    someFiles = []
    for d in range( 1, maxDepth+1 ):
        maxGlob = "/".join( "*" * d )
        topGlob = os.path.join( top, maxGlob )
        allFiles = glob.glob( topGlob )
        someFiles.extend( [ f for f in allFiles if fnmatch.fnmatch( os.path.basename( f ), fnMask ) ] )
    return someFiles

I guess I could also make it a generator with something like this:

def fileNamesRetrieve( top, maxDepth, fnMask  ):
    for d in range( 1, maxDepth+1 ):
        maxGlob = "/".join( "*" * d )
        topGlob = os.path.join( top, maxGlob )
        allFiles = glob.glob( topGlob )
        if fnmatch.fnmatch( os.path.basename( f ), fnMask ):
            yield f

Critique welcome.