feat: add pack extensions and walk method

beet/library/base.py

@@ -29,12 +29,14 @@
     DefaultDict,
     Dict,
     Generic,
+    Iterable,
     Iterator,
     List,
     Mapping,
     MutableMapping,
     Optional,
     Protocol,
+    Set,
     Tuple,
     Type,
     TypeVar,
@@ -128,6 +130,25 @@ def bind(self, namespace: "Namespace", file_type: Type[NamespaceFileType]):
         for key, value in self.items():
             self.process(key, value)
 
+    def generate_tree(self, path: str = "") -> Dict[Any, Any]:
+        """Generate a hierarchy of nested dictionaries representing the files and folders."""
+        prefix = path.split("/") if path else []
+        tree: Dict[Any, Any] = {}
+
+        for filename, file_instance in self.items():
+            parts = filename.split("/")
+
+            if parts[: len(prefix)] != prefix:
+                continue
+
+            parent = tree
+            for part in parts[len(prefix) :]:
+                parent = parent.setdefault(part, {})
+
+            parent[self.file_type] = file_instance
+
+        return tree
+
 
 class NamespacePin(Pin[Type[NamespaceFileType], NamespaceContainer[NamespaceFileType]]):
     """Descriptor for accessing namespace containers by attribute lookup."""
@@ -265,7 +286,12 @@ def get_extra_info(cls) -> Dict[str, Type[PackFile]]:
         return {}
 
     @classmethod
-    def scan(cls, pack: FileOrigin) -> Iterator[Tuple[str, "Namespace"]]:
+    def scan(
+        cls,
+        pack: FileOrigin,
+        extend_namespace: Iterable[Type[NamespaceFile]] = (),
+        extend_namespace_extra: Optional[Mapping[str, Type[PackFile]]] = None,
+    ) -> Iterator[Tuple[str, "Namespace"]]:
         """Load namespaces by walking through a zipfile or directory."""
         name, namespace = None, None
         filenames = (
@@ -275,6 +301,12 @@ def scan(cls, pack: FileOrigin) -> Iterator[Tuple[str, "Namespace"]]:
         )
 
         extra_info = cls.get_extra_info()
+        if extend_namespace_extra:
+            extra_info.update(extend_namespace_extra)
+
+        scope_map = dict(cls.scope_map)
+        for file_type in extend_namespace:
+            scope_map[file_type.scope, file_type.extension] = file_type
 
         for filename in sorted(filenames):
             try:
@@ -297,7 +329,7 @@ def scan(cls, pack: FileOrigin) -> Iterator[Tuple[str, "Namespace"]]:
                 continue
 
             while path := tuple(scope):
-                if file_type := cls.scope_map.get((path, extension)):
+                if file_type := scope_map.get((path, extension)):
                     key = "/".join(
                         filename.relative_to(Path(directory, name, *path)).parts
                     )[: -len(extension)]
@@ -336,6 +368,35 @@ def split_key(self, key: str) -> Tuple[str, str]:
     def join_key(self, key1: str, key2: str) -> str:
         return f"{key1}:{key2}"
 
+    def walk(self) -> Iterator[Tuple[str, Set[str], Dict[str, NamespaceFileType]]]:
+        """Walk over the file hierarchy."""
+        for prefix, namespace in self.proxy.items():
+            separator = ":"
+            roots: List[Tuple[str, Dict[Any, Any]]] = [
+                (prefix, namespace[self.proxy_key].generate_tree())  # type: ignore
+            ]
+
+            while roots:
+                prefix, root = roots.pop()
+
+                dirs: Set[str] = set()
+                files: Dict[str, NamespaceFileType] = {}
+
+                for key, value in root.items():
+                    if not isinstance(key, str):
+                        continue
+                    if any(isinstance(name, str) for name in value):
+                        dirs.add(key)
+                    if file_instance := value.get(self.proxy_key, None):
+                        files[key] = file_instance
+
+                yield prefix + separator, dirs, files
+
+                for directory in dirs:
+                    roots.append((prefix + separator + directory, root[directory]))
+
+                separator = "/"
+
 
 @dataclass
 class NamespaceProxyDescriptor(Generic[NamespaceFileType]):
@@ -379,6 +440,10 @@ class Pack(MatchMixin, MergeMixin, Container[str, NamespaceType]):
     description: PackPin[TextComponent] = PackPin("description", default="")
     pack_format: PackPin[int] = PackPin("pack_format", default=0)
 
+    extend_extra: Dict[str, Type[PackFile]]
+    extend_namespace: List[Type[NamespaceFile]]
+    extend_namespace_extra: Dict[str, Type[PackFile]]
+
     namespace_type: ClassVar[Type[NamespaceType]]
     default_name: ClassVar[str]
     latest_pack_format: ClassVar[int]
@@ -396,6 +461,9 @@ def __init__(
         icon: Optional[PngFile] = None,
         description: Optional[str] = None,
         pack_format: Optional[int] = None,
+        extend_extra: Optional[Mapping[str, Type[PackFile]]] = None,
+        extend_namespace: Iterable[Type[NamespaceFile]] = (),
+        extend_namespace_extra: Optional[Mapping[str, Type[PackFile]]] = None,
     ):
         super().__init__()
         self.name = name
@@ -413,6 +481,10 @@ def __init__(
         if pack_format is not None:
             self.pack_format = pack_format
 
+        self.extend_extra = dict(extend_extra or {})
+        self.extend_namespace = list(extend_namespace)
+        self.extend_namespace_extra = dict(extend_namespace_extra or {})
+
         self.load(path or zipfile)
 
     @overload
@@ -524,8 +596,18 @@ def list_files(
     def get_extra_info(cls) -> Dict[str, Type[PackFile]]:
         return {"pack.mcmeta": JsonFile, "pack.png": PngFile}
 
-    def load(self, origin: Optional[FileOrigin] = None):
+    def load(
+        self,
+        origin: Optional[FileOrigin] = None,
+        extend_extra: Optional[Mapping[str, Type[PackFile]]] = None,
+        extend_namespace: Iterable[Type[NamespaceFile]] = (),
+        extend_namespace_extra: Optional[Mapping[str, Type[PackFile]]] = None,
+    ):
         """Load pack from a zipfile or from the filesystem."""
+        self.extend_extra.update(extend_extra or {})
+        self.extend_namespace.extend(extend_namespace)
+        self.extend_namespace_extra.update(extend_namespace_extra or {})
+
         if origin:
             if not isinstance(origin, ZipFile):
                 origin = Path(origin).resolve()
@@ -546,16 +628,25 @@ def load(self, origin: Optional[FileOrigin] = None):
                 self.name = self.name[:-4]
 
         if origin:
+            extra_info = self.get_extra_info()
+            if self.extend_extra:
+                extra_info.update(self.extend_extra)
+
             files = {
                 filename: loaded
-                for filename, file_type in self.get_extra_info().items()
+                for filename, file_type in extra_info.items()
                 if (loaded := file_type.try_load(origin, filename))
             }
 
             self.extra.merge(files)
 
             namespaces = {
-                name: namespace for name, namespace in self.namespace_type.scan(origin)
+                name: namespace
+                for name, namespace in self.namespace_type.scan(
+                    origin,
+                    self.extend_namespace,
+                    self.extend_namespace_extra,
+                )
             }
 
             self.merge(namespaces)

examples/load_extend/beet.yml

@@ -0,0 +1,4 @@
+require: ["demo.extend_data_pack"]
+data_pack:
+  load: ["src"]
+pipeline: ["demo.process_functions"]

examples/load_extend/demo.py

@@ -0,0 +1,43 @@
+from typing import cast
+
+from beet import Context, JsonFile, NamespaceFile, TextFile, YamlFile
+from beet.library.data_pack import Function
+
+
+class FunctionConfig(YamlFile, NamespaceFile):
+    scope = ("functions",)
+    extension = ".yml"
+
+
+def extend_data_pack(ctx: Context):
+    ctx.data.extend_extra["myproject.json"] = JsonFile
+    ctx.data.extend_namespace.append(FunctionConfig)
+    ctx.data.extend_namespace_extra["numbers.txt"] = TextFile
+
+
+def process_functions(ctx: Context):
+    project_data = cast(JsonFile, ctx.data.extra["myproject.json"]).data
+
+    for prefix, dirs, functions in ctx.data.functions.walk():
+        dirs.discard("zprivate")
+
+        namespace = ctx.data[prefix.partition(":")[0]]
+        numbers = cast(TextFile, namespace.extra["numbers.txt"]).text.splitlines()
+
+        folder_config = ctx.data[FunctionConfig][prefix + "config"].data
+
+        for function in functions.values():
+            function.prepend(
+                Function(
+                    [
+                        f"# config.yml = {folder_config}",
+                        f"# numbers.txt = {numbers}",
+                        f"# myproject.json = {project_data}",
+                    ]
+                )
+            )
+
+    del ctx.data.extra["myproject.json"]
+    for namespace in ctx.data.values():
+        del namespace.extra["numbers.txt"]
+        del namespace[FunctionConfig]

examples/load_extend/src/data/demo/functions/abc/blah.mcfunction

@@ -0,0 +1 @@
+say abc

examples/load_extend/src/data/demo/functions/abc/config.yml

@@ -0,0 +1 @@
+value: 13

examples/load_extend/src/data/demo/functions/config.yml

@@ -0,0 +1 @@
+value: 42

examples/load_extend/src/data/demo/functions/foo.mcfunction

@@ -0,0 +1 @@
+say hello

examples/load_extend/src/data/demo/functions/zprivate/thing.mcfunction

@@ -0,0 +1 @@
+say keep this untouched

examples/load_extend/src/data/demo/numbers.txt

@@ -0,0 +1,5 @@
+1
+2
+3
+4
+5

examples/load_extend/src/data/other/functions/config.yml

@@ -0,0 +1 @@
+value: 7798

examples/load_extend/src/data/other/functions/qqq.mcfunction

@@ -0,0 +1 @@
+say wat

examples/load_extend/src/data/other/numbers.txt

@@ -0,0 +1,3 @@
+1
+2
+9000

examples/load_extend/src/myproject.json

@@ -0,0 +1,3 @@
+{
+  "foo": "bar"
+}

tests/snapshots/examples__build_load_extend__0.data_pack/data/demo/functions/abc/blah.mcfunction

@@ -0,0 +1,4 @@
+# config.yml = {'value': 13}
+# numbers.txt = ['1', '2', '3', '4', '5']
+# myproject.json = {'foo': 'bar'}
+say abc

tests/snapshots/examples__build_load_extend__0.data_pack/data/demo/functions/foo.mcfunction

@@ -0,0 +1,4 @@
+# config.yml = {'value': 42}
+# numbers.txt = ['1', '2', '3', '4', '5']
+# myproject.json = {'foo': 'bar'}
+say hello

tests/snapshots/examples__build_load_extend__0.data_pack/data/demo/functions/zprivate/thing.mcfunction

@@ -0,0 +1 @@
+say keep this untouched

tests/snapshots/examples__build_load_extend__0.data_pack/data/other/functions/qqq.mcfunction

@@ -0,0 +1,4 @@
+# config.yml = {'value': 7798}
+# numbers.txt = ['1', '2', '9000']
+# myproject.json = {'foo': 'bar'}
+say wat

tests/snapshots/examples__build_load_extend__0.data_pack/pack.mcmeta

@@ -0,0 +1,6 @@
+{
+  "pack": {
+    "pack_format": 7,
+    "description": ""
+  }
+}

tests/snapshots/examples__build_load_extend__1.resource_pack/pack.mcmeta

@@ -0,0 +1,6 @@
+{
+  "pack": {
+    "pack_format": 7,
+    "description": ""
+  }
+}