fix: #93 修复目录深度设置后导致目录下的歌曲无法加到播放列表里的问题

xiaomusic/utils.py

@@ -101,63 +101,48 @@ def custom_sort_key(s):
         return (2, s)
 
 
-# fork from https://gist.github.com/dougthor42/001355248518bc64d2f8
+def _get_depth_path(root, directory, depth):
-def walk_to_depth(root, depth=None, *args, **kwargs):
+    # 计算当前目录的深度
-    """
+    relative_path = root[len(directory) :].strip(os.sep)
-    Wrapper around os.walk that stops after going down `depth` folders.
+    path_parts = relative_path.split(os.sep)
-    I had my own version, but it wasn't as efficient as
+    if len(path_parts) >= depth:
-    http://stackoverflow.com/a/234329/1354930, so I modified to be more
+        return os.path.join(directory, *path_parts[:depth])
-    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 root
-            return
-    except IndexError:
-        pass
 
-    try:
+
-        if kwargs["topdown"] is False:
+def _append_files_result(result, root, joinpath, files, support_extension):
-            yield from fallback_func(root, depth, *args, **kwargs)
+    dir_name = os.path.basename(root)
-            return
+    if dir_name not in result:
+        result[dir_name] = []
+    for file in files:
+        # 过滤隐藏文件
+        if file.startswith("."):
+            continue
+        # 过滤文件后缀
+        (name, extension) = os.path.splitext(file)
+        if extension.lower() not in support_extension:
+            continue
+
+        result[dir_name].append(os.path.join(joinpath, file))
+
+
+def traverse_music_directory(
+    directory, depth=10, exclude_dirs=None, support_extension=None
+):
+    result = {}
+    for root, dirs, files in os.walk(directory):
+        # 忽略排除的目录
+        dirs[:] = [d for d in dirs if d not in exclude_dirs]
+
+        # 计算当前目录的深度
+        current_depth = root[len(directory) :].count(os.sep) + 1
+        if current_depth > depth:
+            depth_path = _get_depth_path(root, directory, depth - 1)
+            _append_files_result(result, depth_path, root, files, support_extension)
         else:
-            yield from main_func(root, depth, *args, **kwargs)
+            _append_files_result(result, root, root, files, support_extension)
-            return
+    return result
-    except KeyError:
-        yield from main_func(root, depth, *args, **kwargs)
-        return
 
 
 def downloadfile(url):

xiaomusic/xiaomusic.py

@@ -37,7 +37,7 @@ from xiaomusic.utils import (
     get_local_music_duration,
     get_web_music_duration,
     parse_cookie_string,
-    walk_to_depth,
+    traverse_music_directory,
 )
 
 EOF = object()
@@ -481,39 +481,34 @@ class XiaoMusic:
         encoded_name = urllib.parse.quote(filename)
         return f"http://{self.hostname}:{self.public_port}/{encoded_name}"
 
-    # 递归获取目录下所有歌曲,生成随机播放列表
+    # 获取目录下所有歌曲,生成随机播放列表
     def _gen_all_music_list(self):
         self._all_music = {}
         all_music_by_dir = {}
-        for root, dirs, filenames in walk_to_depth(
+        local_musics = traverse_music_directory(
-            self.music_path, depth=self.music_path_depth
+            self.music_path,
-        ):
+            depth=self.music_path_depth,
-            dirs[:] = [d for d in dirs if d not in self.exclude_dirs]
+            exclude_dirs=self.exclude_dirs,
-            self.log.debug("root:%s dirs:%s music_path:%s", root, dirs, self.music_path)
+            support_extension=SUPPORT_MUSIC_TYPE,
-            dir_name = os.path.basename(root)
+        )
-            if self.music_path == root:
+        for dir_name, files in local_musics.items():
+            if len(files) == 0:
+                continue
+            if dir_name == os.path.basename(self.music_path):
                 dir_name = "其他"
+            if self.music_path != self.download_path and dir_name == os.path.basename(
+                self.download_path
+            ):
+                dir_name = "下载"
             if dir_name not in all_music_by_dir:
                 all_music_by_dir[dir_name] = {}
-            for filename in filenames:
+            for file in files:
-                self.log.debug("gen_all_music_list. filename:%s", filename)
-                # 过滤隐藏文件
-                if filename.startswith("."):
-                    continue
-                # 过滤非音乐文件
-                (name, extension) = os.path.splitext(filename)
-                self.log.debug(
-                    "gen_all_music_list. filename:%s, name:%s, extension:%s",
-                    filename,
-                    name,
-                    extension,
-                )
-                if extension.lower() not in SUPPORT_MUSIC_TYPE:
-                    continue
-
                 # 歌曲名字相同会覆盖
-                self._all_music[name] = os.path.join(root, filename)
+                filename = os.path.basename(file)
+                (name, _) = os.path.splitext(filename)
+                self._all_music[name] = file
                 all_music_by_dir[dir_name][name] = True
+
         self._play_list = list(self._all_music.keys())
         self._cur_play_list = "全部"
         self._gen_play_list()