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.walkstraight out of the source и вставить свой собственный параметр контроля глубины.

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)