Python 3 - дерево каталогов путешествий с ограниченной глубиной рекурсии

Question

Мне нужно рекурсивно обрабатывать все файлы в дереве каталогов, но с ограниченной глубиной.

Это означает, например, чтобы искать файлы в текущем каталоге и первых двух уровнях подкаталога, но не дальше. В этом случае я должен обработать, например, ./subdir1/subdir2/file, но нет ./subdir1/subdir2/subdir3/file,

Как бы я сделал это лучше всего в Python 3?

В настоящее время я использую os.walk обрабатывать все файлы до бесконечной глубины в цикле следующим образом:

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

Я мог бы придумать способ подсчета разделителей каталогов (/) в root определить уровень иерархии текущего файла и break цикл, если этот уровень превышает желаемый максимум.

Я считаю такой подход небезопасным и, возможно, довольно неэффективным, когда нужно игнорировать большое количество подкаталогов. Каков был бы оптимальный подход здесь?

Answer 1

Я думаю, что самый простой и стабильный подход - это скопировать функциональность os.walk прямо из источника и вставьте свой собственный параметр контроля глубины.

import os
import os.path as path

def walk(top, topdown=True, onerror=None, followlinks=False, maxdepth=None):
    islink, join, isdir = path.islink, path.join, path.isdir

    try:
        names = os.listdir(top)
    except OSError, err:
        if onerror is not None:
            onerror(err)
        return

    dirs, nondirs = [], []
    for name in names:
        if isdir(join(top, name)):
            dirs.append(name)
        else:
            nondirs.append(name)

    if topdown:
        yield top, dirs, nondirs

    if maxdepth is None or maxdepth > 1:
        for name in dirs:
            new_path = join(top, name)
            if followlinks or not islink(new_path):
                for x in walk(new_path, topdown, onerror, followlinks, None if maxdepth is None else maxdepth-1):
                    yield x
    if not topdown:
        yield top, dirs, nondirs

for root, dirnames, filenames in walk(args.directory, maxdepth=2):
    #...

---

Если вас не интересуют все эти необязательные параметры, вы можете существенно сократить функцию:

import os

def walk(top, maxdepth):
    dirs, nondirs = [], []
    for name in os.listdir(top):
        (dirs if os.path.isdir(os.path.join(top, name)) else nondirs).append(name)
    yield top, dirs, nondirs
    if maxdepth > 1:
        for name in dirs:
            for x in walk(os.path.join(top, name), maxdepth-1):
                yield x

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

Answer 2

Начиная с python 3.5, os.scandir используется в os.walk вместо os.listdir. Работает во много раз быстрее. Я немного исправил образец @kevin.

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 walkMaxDepth(path, maxdepth-1):
                yield x

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