Travel directory tree with limited recursion depth

前端 未结 2 680
天命终不由人
天命终不由人 2020-11-29 08:59

I need to process all files in a directory tree recursively, but with a limited depth.

That means for example to look for files in the current directory and the fir

2条回答
  •  时光取名叫无心
    2020-11-29 09:28

    I think the easiest and most stable approach would be to copy the functionality of os.walk straight out of the source and insert your own depth-controlling parameter.

    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):
        #...
    

    If you're not interested in all those optional parameters, you can pare down the function pretty substantially:

    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)
    

提交回复
热议问题