diff --git a/xiaomusic/config.py b/xiaomusic/config.py index ae9e1a0..95f37f5 100644 --- a/xiaomusic/config.py +++ b/xiaomusic/config.py @@ -97,6 +97,7 @@ class Config: ffmpeg_location: str = os.getenv("XIAOMUSIC_FFMPEG_LOCATION", "./ffmpeg/bin") active_cmd: str = os.getenv("XIAOMUSIC_ACTIVE_CMD", "play,random_play") exclude_dirs: str = os.getenv("XIAOMUSIC_EXCLUDE_DIRS", "@eaDir") + music_path_depth: int = int(os.getenv("XIAOMUSIC_MUSIC_PATH_DEPTH", "10")) def __post_init__(self) -> None: if self.proxy: diff --git a/xiaomusic/utils.py b/xiaomusic/utils.py index d80ff95..35e424a 100644 --- a/xiaomusic/utils.py +++ b/xiaomusic/utils.py @@ -2,6 +2,7 @@ from __future__ import annotations import difflib +import os import re from collections.abc import AsyncIterator from http.cookies import SimpleCookie @@ -84,3 +85,62 @@ def custom_sort_key(s): else: # 如果前缀和后缀都不是数字,按字典序排序 return (2, s) + + +# fork from https://gist.github.com/dougthor42/001355248518bc64d2f8 +def walk_to_depth(root, depth=None, *args, **kwargs): + """ + Wrapper around os.walk that stops after going down `depth` folders. + I had my own version, but it wasn't as efficient as + http://stackoverflow.com/a/234329/1354930, so I modified to be more + similar to nosklo's answer. + However, nosklo's answer doesn't work if topdown=False, so I kept my + version. + """ + # Let people use this as a standard `os.walk` function. + if depth is None: + return os.walk(root, *args, **kwargs) + + # remove any trailing separators so that our counts are correct. + root = root.rstrip(os.path.sep) + + def main_func(root, depth, *args, **kwargs): + """Faster because it skips traversing dirs that are too deep.""" + root_depth = root.count(os.path.sep) + for dirpath, dirnames, filenames in os.walk(root, *args, **kwargs): + yield (dirpath, dirnames, filenames) + + # calculate how far down we are. + current_folder_depth = dirpath.count(os.path.sep) + if current_folder_depth >= root_depth + depth: + del dirnames[:] + + def fallback_func(root, depth, *args, **kwargs): + """Slower, but works when topdown is False""" + root_depth = root.count(os.path.sep) + for dirpath, dirnames, filenames in os.walk(root, *args, **kwargs): + current_folder_depth = dirpath.count(os.path.sep) + if current_folder_depth <= root_depth + depth: + yield (dirpath, dirnames, filenames) + + # there's gotta be a better way do do this... + try: + if args[0] is False: + yield from fallback_func(root, depth, *args, **kwargs) + return + else: + yield from main_func(root, depth, *args, **kwargs) + return + except IndexError: + pass + + try: + if kwargs["topdown"] is False: + yield from fallback_func(root, depth, *args, **kwargs) + return + else: + yield from main_func(root, depth, *args, **kwargs) + return + except KeyError: + yield from main_func(root, depth, *args, **kwargs) + return diff --git a/xiaomusic/xiaomusic.py b/xiaomusic/xiaomusic.py index c6fbb60..63b7528 100644 --- a/xiaomusic/xiaomusic.py +++ b/xiaomusic/xiaomusic.py @@ -32,6 +32,7 @@ from xiaomusic.utils import ( custom_sort_key, fuzzyfinder, parse_cookie_string, + walk_to_depth, ) EOF = object() @@ -68,6 +69,7 @@ class XiaoMusic: self.ffmpeg_location = config.ffmpeg_location self.active_cmd = config.active_cmd.split(",") self.exclude_dirs = set(config.exclude_dirs.split(",")) + self.music_path_depth = config.music_path_depth # 下载对象 self.download_proc = None @@ -354,7 +356,9 @@ class XiaoMusic: def _gen_all_music_list(self): self._all_music = {} all_music_by_dir = {} - for root, dirs, filenames in os.walk(self.music_path): + for root, dirs, filenames in walk_to_depth( + self.music_path, depth=self.music_path_depth + ): dirs[:] = [d for d in dirs if d not in self.exclude_dirs] self.log.debug("root:%s dirs:%s music_path:%s", root, dirs, self.music_path) dir_name = os.path.basename(root)