def walklevels(path, levels=1):
    path = path.rstrip(os.path.sep)
    assert os.path.isdir(path)
    num_sep = path.count(os.path.sep)
    for root, dirs, files in os.walk(path):
        yield root, dirs, files
        num_sep_this = root.count(os.path.sep)
        if num_sep + levels <= num_sep_this:
            del dirs[:]
Wrap os.walk to allow setting recursion depth.
dataset/ETHPy150Open mikar/demimove/demimove/helpers.py/walklevels