feat: fetch external objects using asset index

Author: vberlier

beet/contrib/vanilla.py

@@ -7,12 +7,14 @@
     "ReleaseRegistry",
     "Release",
     "ClientJar",
+    "AssetIndex",
     "MANIFEST_URL",
+    "RESOURCES_URL",
 ]
 
 
 from pathlib import Path
-from typing import Optional, Union
+from typing import Iterator, Optional, Union
 from zipfile import ZipFile
 
 from pydantic import BaseModel
@@ -25,10 +27,12 @@
     DataPack,
     JsonFile,
     ResourcePack,
+    UnveilMapping,
 )
 from beet.core.utils import FileSystemPath, log_time
 
 MANIFEST_URL: str = "https://piston-meta.mojang.com/mc/game/version_manifest_v2.json"
+RESOURCES_URL: str = "https://resources.download.minecraft.net"
 
 
 class VanillaOptions(BaseModel):
@@ -50,11 +54,15 @@ def __init__(self, cache: Cache, path: FileSystemPath):
         self.assets = ResourcePack()
         self.data = DataPack()
 
-    def mount(self, prefix: Optional[str] = None) -> "ClientJar":
+    def mount(
+        self,
+        prefix: Optional[str] = None,
+        object_mapping: Optional[UnveilMapping] = None,
+    ) -> "ClientJar":
         """Mount the specified prefix if it's not available already."""
         if not prefix:
-            self.mount("assets")
-            self.mount("data")
+            self.mount("assets", object_mapping)
+            self.mount("data", object_mapping)
             return self
 
         if prefix.startswith("assets"):
@@ -70,28 +78,69 @@ def mount(self, prefix: Optional[str] = None) -> "ClientJar":
             with log_time("Extract vanilla pack."):
                 pack.load(ZipFile(self.path))
                 pack.save(path=path)
-                return self
-
-        if pack.path == path.parent:
-            return self
+        elif pack.path != path.parent:
+            pack.unveil(prefix, path)
 
-        pack.unveil(prefix, path)
+        if object_mapping and isinstance(pack, ResourcePack):
+            with self.cache.parallel_downloads():
+                pack.unveil(prefix, object_mapping)
 
         return self
 
 
+class AssetIndex(Container[str, FileSystemPath]):
+    """Class for retrieving assets referenced by a particular release."""
+
+    cache: Cache
+    info: JsonFile
+
+    def __init__(self, cache: Cache, info: JsonFile):
+        super().__init__()
+        self.cache = cache
+        self.info = info
+
+    def missing(self, key: str) -> FileSystemPath:
+        if not key.startswith("assets/"):
+            raise KeyError(key)
+
+        try:
+            object_hash: str = self.info.data["objects"][key[7:]]["hash"]
+        except KeyError as exc:
+            raise KeyError(key) from exc
+
+        path = self.cache.directory / "objects" / object_hash[:2] / object_hash
+
+        if not path.is_file():
+            path.parent.mkdir(parents=True, exist_ok=True)
+            self.cache.download(
+                f"{RESOURCES_URL}/{object_hash[:2]}/{object_hash}",
+                path,
+            )
+
+        return path
+
+    def __iter__(self) -> Iterator[str]:
+        for key in self.info.data["objects"]:
+            yield f"assets/{key}"
+
+    def __len__(self) -> int:
+        return len(self.info.data["objects"])
+
+
 class Release:
     """Class holding information about a minecraft release."""
 
     cache: Cache
     info: JsonFile
 
     _client_jar: Optional[ClientJar]
+    _object_mapping: Optional[UnveilMapping]
 
     def __init__(self, cache: Cache, info: JsonFile):
         self.cache = cache
         self.info = info
         self._client_jar = None
+        self._object_mapping = None
 
     @property
     def type(self) -> str:
@@ -104,8 +153,24 @@ def client_jar(self) -> ClientJar:
             self._client_jar = ClientJar(self.cache, path)
         return self._client_jar
 
-    def mount(self, prefix: Optional[str] = None) -> ClientJar:
-        return self.client_jar.mount(prefix)
+    @property
+    def object_mapping(self) -> UnveilMapping:
+        if not self._object_mapping:
+            path = self.cache.download(self.info.data["assetIndex"]["url"])
+            self._object_mapping = UnveilMapping(
+                AssetIndex(self.cache, JsonFile(source_path=path))
+            )
+        return self._object_mapping
+
+    def mount(
+        self,
+        prefix: Optional[str] = None,
+        fetch_objects: bool = False,
+    ) -> ClientJar:
+        return self.client_jar.mount(
+            prefix=prefix,
+            object_mapping=self.object_mapping if fetch_objects else None,
+        )
 
     @property
     def assets(self) -> ResourcePack:
@@ -182,8 +247,12 @@ def __init__(
         else:
             self.minecraft_version = LATEST_MINECRAFT_VERSION
 
-    def mount(self, prefix: Optional[str] = None) -> ClientJar:
-        return self.releases[self.minecraft_version].mount(prefix)
+    def mount(
+        self,
+        prefix: Optional[str] = None,
+        fetch_objects: bool = False,
+    ) -> ClientJar:
+        return self.releases[self.minecraft_version].mount(prefix, fetch_objects)
 
     @property
     def assets(self) -> ResourcePack:

beet/core/cache.py

@@ -3,17 +3,19 @@
     "Cache",
     "CachePin",
     "CacheTransaction",
+    "DownloadManager",
 ]
 
 
 import json
 import logging
 import shutil
-from contextlib import contextmanager
+from concurrent.futures import Executor, ThreadPoolExecutor
+from contextlib import closing, contextmanager
 from datetime import datetime, timedelta
 from pathlib import Path
 from textwrap import indent
-from typing import Any, ClassVar, Iterator, Optional, Set, Type, TypeVar
+from typing import Any, BinaryIO, ClassVar, Iterator, Optional, Set, Type, TypeVar
 from urllib.request import urlopen
 
 from .container import Container, MatchMixin, Pin
@@ -58,6 +60,7 @@ class Cache:
     index_path: Path
     index: JsonDict
     transaction: CacheTransaction
+    download_manager: "DownloadManager"
 
     index_file: ClassVar[str] = "index.json"
 
@@ -75,6 +78,7 @@ def __init__(
             else self.get_initial_index()
         )
         self.transaction = transaction or CacheTransaction()
+        self.download_manager = DownloadManager()
         self.flush()
 
     def get_initial_index(self) -> JsonDict:
@@ -112,15 +116,22 @@ def get_path(self, key: str) -> Path:
 
         return self.directory / path
 
-    def download(self, url: str) -> Path:
+    @contextmanager
+    def parallel_downloads(self, max_workers: Optional[int] = None):
+        """Launch multiple requests at the same time."""
+        with DownloadManager.parallel(max_workers) as parallel_download_manager:
+            previous_manager = self.download_manager
+            self.download_manager = parallel_download_manager
+            try:
+                yield
+            finally:
+                self.download_manager = previous_manager
+
+    def download(self, url: str, path: Optional[FileSystemPath] = None) -> Path:
         """Download and cache a given url."""
-        path = self.get_path(url)
-
-        if not path.is_file():
-            with log_time('Download "%s".', url), urlopen(url) as f:
-                path.write_bytes(f.read())
-
-        return path
+        if not path:
+            path = self.get_path(url)
+        return self.download_manager.download(url, path)
 
     def has_changed(self, *filenames: Optional[FileSystemPath]) -> bool:
         """Return whether any of the given files changed since the last check."""
@@ -358,3 +369,37 @@ def flush(self):
 
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}({str(self.path)!r})"
+
+
+class DownloadManager:
+    """Download manager."""
+
+    executor: Optional[Executor]
+
+    def __init__(self, executor: Optional[Executor] = None):
+        self.executor = executor
+
+    @classmethod
+    @contextmanager
+    def parallel(cls, max_workers: Optional[int] = None):
+        """Create a download manager that launches multiple requests at the same time."""
+        with ThreadPoolExecutor(max_workers) as executor:
+            yield cls(executor)
+
+    def download(self, url: str, path: FileSystemPath) -> Path:
+        """Download and cache a given url."""
+        path = Path(path)
+
+        if not path.is_file():
+            fileobj = path.open("wb")
+            if self.executor:
+                self.executor.submit(self.retrieve, url, fileobj)
+            else:
+                self.retrieve(url, fileobj)
+
+        return path
+
+    def retrieve(self, url: str, fileobj: BinaryIO):
+        """Retrieve file from url."""
+        with log_time('Download "%s".', url), closing(fileobj), urlopen(url) as f:
+            fileobj.write(f.read())

beet/core/file.py

@@ -27,7 +27,17 @@
 from copy import deepcopy
 from dataclasses import dataclass, replace
 from pathlib import Path
-from typing import Any, Callable, ClassVar, Generic, Optional, Type, TypeVar, Union
+from typing import (
+    Any,
+    Callable,
+    ClassVar,
+    Generic,
+    Mapping,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+)
 from zipfile import ZipFile
 
 import yaml
@@ -62,7 +72,7 @@ def open_image(*args: Any, **kwargs: Any) -> Any:
 SerializeType = TypeVar("SerializeType", bound=Any)
 FileType = TypeVar("FileType", bound="File[Any, Any]")
 
-FileOrigin = Union[FileSystemPath, ZipFile]
+FileOrigin = Union[FileSystemPath, ZipFile, Mapping[str, FileSystemPath]]
 TextFileContent = Union[ValueType, str, None]
 BinaryFileContent = Union[ValueType, bytes, None]
 
@@ -242,6 +252,12 @@ def try_load(
                 return cls(cls.from_zip(origin, str(path)))
             except KeyError:
                 return None
+        elif isinstance(origin, Mapping):
+            try:
+                path = "" if path == Path() else str(path)
+                origin, path = origin[path], ""
+            except KeyError:
+                return None
         path = Path(origin, path)
         return cls(source_path=path) if path.is_file() else None
 
@@ -255,6 +271,8 @@ def dump_zip(self, origin: ZipFile, name: str, raw: SerializeType) -> None:
 
     def dump(self, origin: FileOrigin, path: FileSystemPath):
         """Write the file to a zipfile or to the filesystem."""
+        if isinstance(origin, Mapping):
+            raise TypeError(f'Can\'t dump file "{path}" to read-only mapping.')
         if self._content is None:
             if isinstance(origin, ZipFile):
                 origin.write(self.ensure_source_path(), str(path))