amulet-core 2.0a5__cp311-cp311-macosx_10_9_universal2.whl

amulet/level/abc/_chunk_handle.py

@@ -0,0 +1,335 @@
+from __future__ import annotations
+
+import pickle
+from typing import Optional, TYPE_CHECKING, Generic, TypeVar, Callable, Self
+from collections.abc import Iterator, Iterable
+from contextlib import contextmanager
+from threading import RLock
+from abc import ABC, abstractmethod
+
+from amulet.utils.shareable_lock import LockNotAcquired
+from amulet.chunk import Chunk, get_null_chunk
+from amulet.data_types import DimensionId
+from amulet.errors import ChunkDoesNotExist, ChunkLoadError
+from amulet.utils.signal import Signal
+
+from ._level import LevelFriend, LevelT
+from ._history import HistoryManagerLayer
+
+
+if TYPE_CHECKING:
+    from ._level import Level
+    from ._raw_level import RawDimension
+
+
+ChunkT = TypeVar("ChunkT", bound=Chunk)
+RawDimensionT = TypeVar("RawDimensionT", bound="RawDimension")
+
+
+class ChunkKey(tuple[int, int]):
+    def __new__(cls, cx: int, cz: int) -> Self:
+        return tuple.__new__(cls, (cx, cz))
+
+    def __init__(self, cx: int, cz: int) -> None:
+        self._bytes: Optional[bytes] = None
+
+    @property
+    def cx(self) -> int:
+        return self[0]
+
+    @property
+    def cz(self) -> int:
+        return self[1]
+
+    def __bytes__(self) -> bytes:
+        if self._bytes is None:
+            self._bytes = b"/".join((str(self[0]).encode(), str(self[1]).encode()))
+        return self._bytes
+
+
+class ChunkHandle(
+    LevelFriend[LevelT],
+    ABC,
+    Generic[LevelT, RawDimensionT, ChunkT],
+):
+    """
+    A class which manages chunk data.
+    You must acquire the lock for the chunk before reading or writing data.
+    Some internal synchronisation is done to catch some threading issues.
+    """
+
+    _lock: RLock
+    _dimension: DimensionId
+    _key: ChunkKey
+    _chunk_history: HistoryManagerLayer[ChunkKey]
+    _chunk_data_history: HistoryManagerLayer[bytes]
+    _raw_dimension: Optional[RawDimensionT]
+
+    __slots__ = (
+        "_lock",
+        "_dimension_id",
+        "_key",
+        "_chunk_history",
+        "_chunk_data_history",
+        "_raw_dimension",
+    )
+
+    def __init__(
+        self,
+        level_ref: Callable[[], LevelT | None],
+        chunk_history: HistoryManagerLayer[ChunkKey],
+        chunk_data_history: HistoryManagerLayer[bytes],
+        dimension_id: DimensionId,
+        cx: int,
+        cz: int,
+    ) -> None:
+        super().__init__(level_ref)
+        self._lock = RLock()
+        self._dimension_id = dimension_id
+        self._key = ChunkKey(cx, cz)
+        self._chunk_history = chunk_history
+        self._chunk_data_history = chunk_data_history
+        self._raw_dimension = None
+
+    changed = Signal[()]()
+
+    @property
+    def dimension_id(self) -> DimensionId:
+        return self._dimension_id
+
+    @property
+    def cx(self) -> int:
+        return self._key.cx
+
+    @property
+    def cz(self) -> int:
+        return self._key.cz
+
+    def _get_raw_dimension(self) -> RawDimensionT:
+        if self._raw_dimension is None:
+            self._raw_dimension = self._l.raw.get_dimension(self.dimension_id)
+        return self._raw_dimension
+
+    @contextmanager
+    def lock(
+        self,
+        *,
+        blocking: bool = True,
+        timeout: float = -1,
+    ) -> Iterator[None]:
+        """
+        Lock access to the chunk.
+
+        >>> level: Level
+        >>> dimension_name: str
+        >>> cx: int
+        >>> cz: int
+        >>> with level.get_dimension(dimension_name).get_chunk_handle(cx, cz).lock():
+        >>>     # Do what you need to with the chunk
+        >>>     # No other threads are able to edit or set the chunk while in this with block.
+
+        If you want to lock, get and set the chunk data :meth:`edit` is probably a better fit.
+
+        :param blocking: Should this block until the lock is acquired.
+        :param timeout: The amount of time to wait for the lock.
+        :raises:
+            LockNotAcquired: If the lock could not be acquired.
+        """
+        if not self._lock.acquire(blocking, timeout):
+            # Thread was not acquired
+            raise LockNotAcquired("Lock was not acquired.")
+        try:
+            yield
+        finally:
+            self._lock.release()
+
+    @contextmanager
+    def edit(
+        self,
+        *,
+        components: Iterable[str] | None = None,
+        blocking: bool = True,
+        timeout: float = -1,
+    ) -> Iterator[ChunkT | None]:
+        """Lock and edit a chunk.
+
+        If you only want to access/modify parts of the chunk data you can specify the components you want to load.
+        This makes it faster because you don't need to load unneeded parts.
+
+        >>> level: Level
+        >>> dimension_name: str
+        >>> cx: int
+        >>> cz: int
+        >>> with level.get_dimension(dimension_name).get_chunk_handle(cx, cz).edit() as chunk:
+        >>>     # Edit the chunk data
+        >>>     # No other threads are able to edit the chunk while in this with block.
+        >>>     # When the with block exits the edited chunk will be automatically set if no exception occurred.
+
+        :param components: None to load all components or an iterable of component strings to load.
+        :param blocking: Should this block until the lock is acquired.
+        :param timeout: The amount of time to wait for the lock.
+        :raises:
+            LockNotAcquired: If the lock could not be acquired.
+        """
+        with self.lock(blocking=blocking, timeout=timeout):
+            chunk = self.get(components)
+            yield chunk
+            # If an exception occurs in user code, this line won't be run.
+            self._set(chunk)
+        self.changed.emit()
+        self._l.changed.emit()
+
+    def exists(self) -> bool:
+        """
+        Does the chunk exist. This is a quick way to check if the chunk exists without loading it.
+
+        This state may change if the lock is not acquired.
+
+        :return: True if the chunk exists. Calling get on this chunk handle may still throw ChunkLoadError
+        """
+        if self._chunk_history.has_resource(self._key):
+            return self._chunk_history.resource_exists(self._key)
+        else:
+            # The history system is not aware of the chunk. Look in the level data
+            return self._get_raw_dimension().has_chunk(self.cx, self.cz)
+
+    def _preload(self) -> None:
+        """Load the chunk data if it has not already been loaded."""
+        if not self._chunk_history.has_resource(self._key):
+            # The history system is not aware of the chunk. Load from the level data
+            chunk: Chunk
+            try:
+                raw_chunk = self._get_raw_dimension().get_raw_chunk(self.cx, self.cz)
+                chunk = self._get_raw_dimension().raw_chunk_to_native_chunk(
+                    raw_chunk,
+                    self.cx,
+                    self.cz,
+                )
+            except ChunkDoesNotExist:
+                self._chunk_history.set_initial_resource(self._key, b"")
+            except ChunkLoadError as e:
+                self._chunk_history.set_initial_resource(self._key, pickle.dumps(e))
+            else:
+                self._chunk_history.set_initial_resource(
+                    self._key, pickle.dumps(chunk.chunk_id)
+                )
+                for component_id, component_data in chunk.serialise_chunk().items():
+                    if component_data is None:
+                        raise RuntimeError(
+                            "Component must not be None when initialising chunk"
+                        )
+                    self._chunk_data_history.set_initial_resource(
+                        b"/".join((bytes(self._key), component_id.encode())),
+                        component_data,
+                    )
+
+    def _get_null_chunk(self) -> ChunkT:
+        """Get a null chunk instance used for this chunk.
+
+        :raises:
+            ChunkDoesNotExist if the chunk does not exist.
+        """
+        data = self._chunk_history.get_resource(self._key)
+        if data:
+            obj: ChunkLoadError | str = pickle.loads(data)
+            if isinstance(obj, ChunkLoadError):
+                raise obj
+            elif isinstance(obj, str):
+                return get_null_chunk(obj)  # type: ignore
+            else:
+                raise RuntimeError
+        else:
+            raise ChunkDoesNotExist
+
+    def get_class(self) -> type[ChunkT]:
+        """Get the chunk class used for this chunk.
+
+        :raises:
+            ChunkDoesNotExist if the chunk does not exist.
+        """
+        return type(self._get_null_chunk())
+
+    def get(self, components: Iterable[str] | None = None) -> ChunkT:
+        """Get a unique copy of the chunk data.
+
+        If you want to edit the chunk, use :meth:`edit` instead.
+
+        If you only want to access/modify parts of the chunk data you can specify the components you want to load.
+        This makes it faster because you don't need to load unneeded parts.
+
+        :param components: None to load all components or an iterable of component strings to load.
+        :return: A unique copy of the chunk data.
+        """
+        with self.lock(blocking=False):
+            self._preload()
+            chunk = self._get_null_chunk()
+            if components is None:
+                components = chunk.component_ids
+            else:
+                # Ensure all component ids are valid for this class.
+                components = set(components).intersection(chunk.component_ids)
+            chunk_components = dict[str, bytes | None]()
+            for component_id in components:
+                chunk_components[component_id] = self._chunk_data_history.get_resource(
+                    b"/".join((bytes(self._key), component_id.encode()))
+                )
+            chunk.reconstruct_chunk(chunk_components)
+            return chunk
+
+    def _set(self, chunk: ChunkT | None) -> None:
+        """Lock must be acquired before calling this"""
+        history = self._chunk_history
+        if not history.has_resource(self._key):
+            if self._l.history_enabled:
+                self._preload()
+            else:
+                history.set_initial_resource(self._key, b"")
+        if chunk is None:
+            history.set_resource(self._key, b"")
+        else:
+            self._validate_chunk(chunk)
+            try:
+                old_chunk_class = self.get_class()
+            except ChunkLoadError:
+                old_chunk_class = None
+            new_chunk_class = type(chunk)
+            component_data = chunk.serialise_chunk()
+            if old_chunk_class != new_chunk_class and None in component_data.values():
+                raise RuntimeError(
+                    "When changing chunk class all the data must be present."
+                )
+            history.set_resource(self._key, pickle.dumps(new_chunk_class))
+            for component_id, data in component_data.items():
+                if data is None:
+                    continue
+                self._chunk_data_history.set_resource(
+                    b"/".join((bytes(self._key), component_id.encode())),
+                    data,
+                )
+
+    @staticmethod
+    @abstractmethod
+    def _validate_chunk(chunk: ChunkT) -> None:
+        raise NotImplementedError
+
+    def set(self, chunk: ChunkT) -> None:
+        """
+        Overwrite the chunk data.
+        You must acquire the chunk lock before setting.
+        If you want to edit the chunk, use :meth:`edit` instead.
+
+        :param chunk: The chunk data to set.
+        :raises:
+            LockNotAcquired: If the chunk is already locked by another thread.
+        """
+        with self.lock(blocking=False):
+            self._set(chunk)
+        self.changed.emit()
+        self._l.changed.emit()
+
+    def delete(self) -> None:
+        """Delete the chunk from the level."""
+        with self.lock(blocking=False):
+            self._set(None)
+        self.changed.emit()
+        self._l.changed.emit()

amulet/level/abc/_dimension.py

@@ -0,0 +1,86 @@
+from abc import ABC, abstractmethod
+from typing import Generic, TypeVar, Callable
+from weakref import WeakValueDictionary
+from threading import Lock
+
+from amulet.data_types import DimensionId
+from amulet.block import BlockStack
+from amulet.biome import Biome
+from amulet.selection import SelectionGroup
+
+from ._level import LevelFriend, LevelT
+from ._history import HistoryManagerLayer
+from ._chunk_handle import ChunkKey, ChunkHandle
+from ._raw_level import RawDimension
+
+
+ChunkHandleT = TypeVar("ChunkHandleT", bound=ChunkHandle)
+RawDimensionT = TypeVar("RawDimensionT", bound=RawDimension)
+
+
+class Dimension(LevelFriend[LevelT], ABC, Generic[LevelT, RawDimensionT, ChunkHandleT]):
+    _dimension_id: DimensionId
+    _chunk_handles: WeakValueDictionary[tuple[int, int], ChunkHandleT]
+    _chunk_handle_lock: Lock
+    _chunk_history: HistoryManagerLayer[ChunkKey]
+    _chunk_data_history: HistoryManagerLayer[bytes]
+    _raw: RawDimensionT
+
+    def __init__(
+        self, level_ref: Callable[[], LevelT | None], dimension_id: DimensionId
+    ) -> None:
+        super().__init__(level_ref)
+        self._dimension_id = dimension_id
+        self._chunk_handles = WeakValueDictionary()
+        self._chunk_handle_lock = Lock()
+        self._chunk_history = self._l._o.history_manager.new_layer()
+        self._chunk_data_history = self._l._o.history_manager.new_layer()
+        self._raw = self._l.raw.get_dimension(self._dimension_id)
+
+    @property
+    def dimension_id(self) -> DimensionId:
+        return self._dimension_id
+
+    def bounds(self) -> SelectionGroup:
+        """The editable region of the dimension."""
+        return self._raw.bounds()
+
+    def default_block(self) -> BlockStack:
+        """The default block for this dimension"""
+        return self._raw.default_block()
+
+    def default_biome(self) -> Biome:
+        """The default biome for this dimension"""
+        return self._raw.default_biome()
+
+    def chunk_coords(self) -> set[tuple[int, int]]:
+        """
+        The coordinates of every chunk that exists in this dimension.
+
+        This is the combination of chunks saved to the level and chunks yet to be saved.
+        """
+        chunks: set[tuple[int, int]] = set(self._raw.all_chunk_coords())
+        for key, state in self._chunk_history.resources_exist_map().items():
+            if state:
+                chunks.add((key.cx, key.cz))
+            else:
+                chunks.discard((key.cx, key.cz))
+        return chunks
+
+    def changed_chunk_coords(self) -> set[tuple[int, int]]:
+        """The coordinates of every chunk in this dimension that have been changed since the last save."""
+        return {(key.cx, key.cz) for key in self._chunk_history.changed_resources()}
+
+    @abstractmethod
+    def _create_chunk_handle(self, cx: int, cz: int) -> ChunkHandleT:
+        raise NotImplementedError
+
+    def get_chunk_handle(self, cx: int, cz: int) -> ChunkHandleT:
+        key = cx, cz
+        with self._chunk_handle_lock:
+            chunk_handle = self._chunk_handles.get(key)
+            if chunk_handle is None:
+                chunk_handle = self._chunk_handles[key] = self._create_chunk_handle(
+                    cx, cz
+                )
+            return chunk_handle

amulet/level/abc/_history/__init__.py

@@ -0,0 +1 @@
+from ._history_manager import HistoryManager, HistoryManagerLayer

amulet/level/abc/_history/_cache.py

@@ -0,0 +1,224 @@
+from __future__ import annotations
+
+import shutil
+from typing import Optional, Callable, cast, IO, Self
+from threading import Lock
+import os
+from weakref import ref, finalize
+import time
+import glob
+import re
+import tempfile
+
+import portalocker
+from leveldb import LevelDB
+from amulet.utils.weakref import CallableWeakMethod
+
+
+TempPattern = re.compile(r"amulettmp.*?-(?P<time>\d+)")
+
+
+def _temp_dir() -> str:
+    temp_dir = os.environ.get("CACHE_DIR")
+    if temp_dir is None:
+        raise RuntimeError
+    return temp_dir
+
+
+def _clear_temp_dirs() -> None:
+    """
+    Try and delete historic temporary directories.
+    If things went very wrong in past sessions temporary directories may still exist.
+    """
+    for path in glob.glob(
+        os.path.join(glob.escape(tempfile.gettempdir()), "amulettmp*")
+    ) + glob.glob(
+        os.path.join(glob.escape(_temp_dir()), "**", "amulettmp*"), recursive=True
+    ):
+        name = os.path.basename(path)
+        match = TempPattern.fullmatch(name)
+        if match and int(match.group("time")) < (time.time() - 7 * 24 * 3600):
+            lock_path = os.path.join(path, "lock")
+            if os.path.exists(lock_path):
+                with open(lock_path) as lock:
+                    # make sure it is not locked by another process
+                    try:
+                        portalocker.lock(lock, portalocker.LockFlags.EXCLUSIVE)
+                    except:
+                        continue
+                    else:
+                        portalocker.unlock(lock)
+            shutil.rmtree(path, ignore_errors=True)
+
+
+_clear_temp_dirs()
+
+
+class TempDir(str):
+    """
+    A temporary directory to do with as you wish.
+
+    >>> t = TempDir()
+    >>> path = os.path.join(t, "your_file.txt")  # TempDir is a subclass of str
+    >>> # make sure all files in the temporary directory are closed before releasing or closing this object.
+    >>> # The temporary directory will be deleted when the last reference to `t` is lost or when `t.close()` is called
+    """
+
+    __lock: IO | None
+    __finalise: finalize
+
+    def __new__(cls, group: str) -> Self:
+        cache_dir = os.path.join(_temp_dir(), group)
+        os.makedirs(cache_dir, exist_ok=True)
+        return super().__new__(
+            cls,
+            tempfile.mkdtemp(
+                prefix="amulettmp",
+                suffix=f"-{time.time():.0f}",
+                dir=cache_dir,
+            ),
+        )
+
+    def __init__(self, group: str) -> None:
+        self.__lock = open(os.path.join(self, "lock"), "w")
+        portalocker.lock(self.__lock, portalocker.LockFlags.EXCLUSIVE)
+        self.__finalise = finalize(self, CallableWeakMethod(self._close))
+
+    def _close(self) -> None:
+        if self.__lock is not None:
+            portalocker.unlock(self.__lock)
+            self.__lock.close()
+            self.__lock = None
+            shutil.rmtree(self)
+
+    def close(self) -> None:
+        """Close the lock and delete the directory."""
+        self.__finalise()
+
+    def __del__(self) -> None:
+        self.__finalise()
+
+
+class DiskCache:
+    """
+    A key, value database with a fast access RAM component and a longer term storage disk component.
+    Keys and values are both bytes.
+    The disk component is a leveldb database.
+    """
+
+    def __init__(self, path: str, max_size: int) -> None:
+        """
+        Create a new DiskCache
+        :param path: The path to save the disk component to.
+        :param max_size: The maximum amount of RAM that values can occupy. Key size is assumed negligible.
+            When this is overflowed, the least recently used entries are unloaded to the disk storage.
+        """
+        self._lock = Lock()
+        self._ram: dict[bytes, tuple[bytes, bool]] = {}
+        self._path = path
+        self._disk = LevelDB(path, create_if_missing=True)
+        self._max_size: int = max_size
+        self._size: int = 0
+        self.__finalise = finalize(self, CallableWeakMethod(self._close))
+
+    def _close(self) -> None:
+        self._disk.close()
+        shutil.rmtree(self._path, ignore_errors=True)
+
+    def __del__(self) -> None:
+        self.__finalise()
+
+    @property
+    def max_size(self) -> int:
+        return self._max_size
+
+    @max_size.setter
+    def max_size(self, max_size: int) -> None:
+        if not isinstance(max_size, int):
+            raise TypeError
+        with self._lock:
+            self._max_size = max_size
+            self._free()
+
+    def __setitem__(self, key: bytes, value: bytes) -> None:
+        with self._lock:
+            self._remove(key)
+            self._ram[key] = (value, True)
+            self._size += len(value)
+            self._free()
+
+    def _remove(self, key: bytes) -> None:
+        if key in self._ram:
+            data = self._ram.pop(key)[0]
+            self._size -= len(data)
+
+    def __delitem__(self, key: bytes) -> None:
+        with self._lock:
+            self._remove(key)
+            if key in self._disk:
+                del self._disk[key]
+
+    def _free(self) -> None:
+        """Push some values to disk"""
+        if self._size > self._max_size:
+            keys = iter(self._ram.copy())
+            while self._size > self._max_size:
+                key = next(keys)
+                value, changed = self._ram.pop(key)
+                self._size -= len(value)
+                if changed:
+                    self._disk[key] = value
+
+    def __getitem__(self, key: bytes) -> bytes:
+        with self._lock:
+            if key in self._ram:
+                value = self._ram.pop(key)
+                # Push it to the end
+                self._ram[key] = value
+                return value[0]
+            elif key in self._disk:
+                data = self._disk[key]
+                self._ram[key] = (data, False)
+                self._size += len(data)
+                self._free()
+                return data
+            else:
+                raise KeyError
+
+
+class GlobalDiskCache(DiskCache):
+    _instance_ref: Callable[[], Optional[GlobalDiskCache]] = cast(
+        Callable[[], Optional["GlobalDiskCache"]], lambda: None
+    )
+    _cache_size = 100_000_000
+
+    @classmethod
+    def instance(cls) -> GlobalDiskCache:
+        """
+        Get the global disk cache instance.
+        The caller must store a strong reference to the returned value otherwise it will be destroyed.
+        """
+        instance: Optional[GlobalDiskCache] = cls._instance_ref()
+        if instance is None:
+            instance = GlobalDiskCache()
+            cls._instance_ref = ref(instance)
+        return instance
+
+    @classmethod
+    def cache_size(cls) -> int:
+        instance: Optional[GlobalDiskCache] = cls._instance_ref()
+        if instance is None:
+            return cls._cache_size
+        else:
+            return instance.max_size
+
+    @classmethod
+    def set_cache_size(cls, size: int) -> None:
+        instance: Optional[GlobalDiskCache] = cls._instance_ref()
+        cls._cache_size = size
+        if instance is not None:
+            instance.max_size = size
+
+    def __init__(self) -> None:
+        self._temp_dir = TempDir("level_data")
+        super().__init__(os.path.join(self._temp_dir, "history_db"), 100_000_000)