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 the code started to look ugly when counting the depth etc.

What is really the most "neat" implementation?

Answer 1

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 2

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 3

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.

Answer 4

Here's a simple function to do this

import os
from glob import glob
from pathlib import Path

def find_sub_dirs(path, depth=2):
    path = Path(path)
    assert path.exists(), f'Path: {path} does not exist'
    depth_search = '*/' * depth
    search_pattern = os.path.join(path, depth_search)
    return list(glob(f'{search_pattern}'))

Answer 5

Note that @phihag's great suggestion also works with pathlib.Path.glob:

from pathlib import Path
from typing import List

def get_subitems(folder: Path, level: int) -> List[Path]:
    if level == 0:
        return [item]
    pattern = "/".join(["*"] * level)
    return sorted(item for item in item.glob(pattern))