Skip to content

Instantly share code, notes, and snippets.

@TheMatt2
Last active November 14, 2024 05:29
Show Gist options
  • Save TheMatt2/faf5ca760c61a267412c46bb977718fa to your computer and use it in GitHub Desktop.
Save TheMatt2/faf5ca760c61a267412c46bb977718fa to your computer and use it in GitHub Desktop.
A python function to do an os.walk(), but only to a certain depth. A negative depth indicates full depth.
# MIT License
#
# Copyright (c) 2021 Matthew Schweiss
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
# Partially from https://stackoverflow.com/questions/229186/os-walk-without-digging-into-directories-below
import os
def walklevel(path, depth = 1):
"""It works just like os.walk, but you can pass it a level parameter
that indicates how deep the recursion will go.
If depth is 1, the current directory is listed.
If depth is 0, nothing is returned.
If depth is -1 (or less than 0), the full depth is walked.
"""
# If depth is negative, just walk
# Not using yield from for python2 compat
# and copy dirs to keep consistant behavior for depth = -1 and depth = inf
if depth < 0:
for root, dirs, files in os.walk(path):
yield root, dirs[:], files
return
elif depth == 0:
return
# path.count(os.path.sep) is safe because
# - On Windows "\\" is never allowed in the name of a file or directory
# - On UNIX "/" is never allowed in the name of a file or directory
# - On MacOS a literal "/" is quitely translated to a ":" so it is still
# safe to count "/".
base_depth = path.rstrip(os.path.sep).count(os.path.sep)
for root, dirs, files in os.walk(path):
yield root, dirs[:], files
cur_depth = root.count(os.path.sep)
if base_depth + depth <= cur_depth:
del dirs[:]
@apfelchips
Copy link

Thanks. I had a use for it here:
randy3k/ProjectManager@f45065f

@mikecando
Copy link

Just before line 26 base_depth = path.rstrip(os.path.sep).count(os.path.sep)

I needed to add the line
depth -= 1

It seems when depth = 0 for that code, I get the current directory results. I am running this on windows. This is a hack for sure but seems to work perfectly. If some one has a better suggestion, please share. Thank you so much. This is awesome

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment