仿照 os.walk 寫了一個 generator lwalk, 他的行為如同 os.walk 但是多了一個 max_level 可以控制最大的遍歷深度, 為了與 os.walk 盡量吻合, 我也實作了 topdown 和 followlinks 這兩個 arguments, 但為了不使 code 太過複雜, 我省略了 onerror 參數的實作以及若干 error handling。

代碼如下:

import os

def lwalk(top, topdown=True, followlinks=False, max_level=None):

if max_level is None:

new_max_level = None

else:

if max_level==0:

return

else:

new_max_level = max_level-1

top = os.fspath(top)

dirs, nondirs, walk_dirs = [], [], []

with os.scandir(top) as it:

for entry in it:

if entry.is_dir():

dirs.append(entry.name)

else:

nondirs.append(entry.name)

if not topdown and entry.is_dir():

if followlinks or not entry.is_symlink():

walk_dirs.append(entry.path)

if topdown:

yield top, dirs, nondirs

for dirname in dirs:

new_path = os.path.join(top, dirname)

if followlinks or not os.path.islink(new_path):

yield from lwalk(new_path, topdown, followlinks, new_max_level)

else:

for new_path in walk_dirs:

yield from lwalk(new_path, topdown, followlinks, new_max_level)

yield top, dirs, nondirs

簡單的範例如下:

for root, dirs, files in lwalk('YOUR_TOP_PATH', max_level=4):

print(root, dirs, files)

一些說明:

核心在於使用 os.scandir 來保證系統遍歷的效率

使用 max_level 來控制最大遍歷深度, 在 recursively 進入下一層的時候, 將最大深度減 1

要實作 buttom up, 則需先 recursively 進入下一層再 yield 目錄與文件

這邊有一個省略掉 topdown, followlink 和若干處理的簡單版本, 可以幫助你觀察一下核心的實作手段:

import os

def lwalk(top, max_level=10000):

if max_level==0:

return

dirs, nondirs = [], []

with os.scandir(top) as it:

for entry in it:

if entry.is_dir():

dirs.append(entry.name)

else:

nondirs.append(entry.name)

yield top, dirs, nondirs

for dirname in dirs:

new_path = os.path.join(top, dirname)

yield from lwalk(new_path, max_level-1)

for root, dirs, files in lwalk('a', max_level=4):

print(root, dirs, files)