fix: support overlays in select_files

mcbeet · Oct 7, 2023 · 1297472 · 1297472
1 parent 08be63f
commit 1297472
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 7 deletions.
diff --git a/beet/toolchain/select.py b/beet/toolchain/select.py
@@ -15,6 +15,7 @@
 from typing import (
     Any,
     Dict,
+    Iterable,
     Literal,
     Optional,
     Tuple,
@@ -30,7 +31,7 @@
 
 from beet.core.file import File
 from beet.core.utils import snake_case
-from beet.library.base import Pack
+from beet.library.base import NamespaceFile, Pack
 
 from .config import ListOption
 from .template import TemplateManager
@@ -199,15 +200,11 @@ def select_files(
 
                 for group, spec in self.match_spec.items():
                     if file_type := group_map.get(group):
-                        for path, file_instance in pack[file_type].items():
-                            if spec.match_file(path):
-                                result[file_instance] = None, path
+                        result.update(_gather_from_pack(pack, file_type, spec))
 
             else:
                 for file_type in file_types:
-                    for path, file_instance in pack[file_type].items():
-                        if self.match_spec.match_file(path):
-                            result[file_instance] = None, path
+                    result.update(_gather_from_pack(pack, file_type, self.match_spec))
 
         return result
 
@@ -249,3 +246,14 @@ def select_files(
         if extend
         else selector.select_files(pack, *extensions)
     )
+
+
+def _gather_from_pack(
+    pack: Pack[Any], file_type: Type[NamespaceFile], spec: PathSpec
+) -> Iterable[Tuple[NamespaceFile, Tuple[Optional[str], Optional[str]]]]:
+    for path, file_instance in pack[file_type].items():
+        if spec.match_file(path):
+            yield file_instance, (None, path)
+    if pack.overlay_parent is None:
+        for overlay in pack.overlays.values():
+            yield from _gather_from_pack(overlay, file_type, spec)
diff --git a/tests/test_data_pack.py b/tests/test_data_pack.py
@@ -694,3 +694,8 @@ def test_overlay():
     assert p.overlays["d"].functions["demo:init"] == Function(["say original init"])
     assert p.overlays["bop"].functions["demo:init"] == Function(["say init"])
     assert p.overlays["bop2"].functions["demo:init"] == Function()
+
+    f1 = set(select_files(p, files=r".*", extend=Function))
+    f2 = set(select_files(p, match="*"))
+    assert f1 == f2
+    assert len(f1) == 6