[Feature Request] Level API restructure

[gentlegiantJGC]

Feature Request

The Problem

The top level API for the level class is a little cluttered. I would like this to be less cluttered.

Feature Description

I think all the chunk attributes should be grouped under the same attribute. Likewise with other objects.

This ties into Amulet-Team/Amulet-Core#260

from __future__ import annotations

from functools import cached_property
from weakref import proxy, WeakValueDictionary
from threading import RLock, Lock
from contextlib import contextmanager
from collections import deque
from copy import deepcopy


class Chunk:
    pass


class LockNotAcquired(RuntimeError):
    pass


class ChunkStorage:
    def __init__(self, level: Level):
        # Weak pointer to the level to get raw and shared data
        self._level: Level = proxy(level)
        # Mapping from chunk location to chunk object. Weakly stored so that we don't need to manually unload.
        self._chunks = WeakValueDictionary[tuple[str, int, int], Chunk]()
        # A deque to keep recently/frequently used chunks loaded
        self._chunk_cache = deque[Chunk](maxlen=100)
        # A lock per chunk
        self._locks = WeakValueDictionary[tuple[str, int, int], RLock]()
        # A lock that must be acquired before touching _locks
        self._locks_lock = Lock()

    def __get_lock(self, key: tuple[str, int, int]) -> RLock:
        with self._locks_lock:
            lock = self._locks.get(key)
            if lock is None:
                lock = self._locks[key] = RLock()
        return lock

    @contextmanager
    def lock(self, dimension: str, cx: int, cz: int, *, blocking: bool = True, timeout: float = -1):
        """
        Lock access to the chunk.

        >>> with level.chunk.lock(dimension, cx, cz):
        >>>     # 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 dimension: The dimension the chunk is stored in.
        :param cx: The chunk x coordinate.
        :param cz: The chunk z coordinate.
        :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 was not acquired.
        """
        key = (dimension, cx, cz)
        lock = self.__get_lock(key)
        if not lock.acquire(blocking, timeout):
            # Thread was not acquired
            raise LockNotAcquired("Lock was not acquired.")
        try:
            yield
        finally:
            lock.release()

    @contextmanager
    def edit(self, dimension: str, cx: int, cz: int, blocking: bool = True, timeout: float = -1):
        """
        Lock and edit a chunk.

        >>> with level.chunk.edit(dimension, cx, cz) 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.
        """
        with self.lock(dimension, cx, cz, blocking=blocking, timeout=timeout):
            chunk = self.get(dimension, cx, cz)
            yield chunk
            # If an exception occurs in user code, this line won't be run.
            self.set(dimension, cx, cz, chunk)

    def get(self, dimension: str, cx: int, cz: int) -> Chunk:
        """
        Get a deep copy of the chunk data.
        If you want to edit the chunk, use :meth:`edit` instead.

        :param dimension: The dimension the chunk is stored in.
        :param cx: The chunk x coordinate.
        :param cz: The chunk z coordinate.
        :return: A unique copy of the chunk data.
        """
        return Chunk()

    def set(self, dimension: str, cx: int, cz: int, chunk: Chunk):
        """
        Overwrite the chunk data.
        You must lock access to the chunk before setting it otherwise an exception may be raised.
        If you want to edit the chunk, use :meth:`edit` instead.

        :param dimension: The dimension the chunk is stored in.
        :param cx: The chunk x coordinate.
        :param cz: The chunk z coordinate.
        :param chunk: The chunk data to set.
        :raises:
            LockNotAcquired: If the chunk is already locked by another thread.
        """
        key = (dimension, cx, cz)
        lock = self.__get_lock(key)
        if lock.acquire(False):
            try:
                chunk = deepcopy(chunk)
                # TODO set the chunk and notify listeners
            finally:
                lock.release()
        else:
            raise LockNotAcquired("Cannot set a chunk if it is locked by another thread.")

    def on_change(self, callback):
        """A notification system for chunk changes."""
        raise NotImplementedError


class Level:
    @cached_property
    def chunk(self) -> ChunkStorage:
        return ChunkStorage(self)