From 91a8c9eb50d910a8a9de3fac5c27aefed8205fcb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=B6=B5=E6=9B=A6?= <im.hanxi@gmail.com>
Date: Sun, 23 Jun 2024 03:30:32 +0000
Subject: [PATCH] =?UTF-8?q?fix:=20#76=20=E6=96=B0=E5=A2=9EXIAOMUSIC=5FMUSI?=
 =?UTF-8?q?C=5FPATH=5FDEPTH=E9=85=8D=E7=BD=AE=E7=94=9F=E6=88=90=E6=92=AD?=
 =?UTF-8?q?=E6=94=BE=E5=88=97=E8=A1=A8=E7=9A=84=E7=9B=AE=E5=BD=95=E6=B7=B1?=
 =?UTF-8?q?=E5=BA=A6=EF=BC=8C=E9=BB=98=E8=AE=A410?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 xiaomusic/config.py    |  1 +
 xiaomusic/utils.py     | 60 ++++++++++++++++++++++++++++++++++++++++++
 xiaomusic/xiaomusic.py |  6 ++++-
 3 files changed, 66 insertions(+), 1 deletion(-)

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)