Added some chunk attribute classes

amulet/chunk/biome.py

@@ -0,0 +1,75 @@
+from typing import Union, Iterable
+
+import numpy
+from numpy.typing import ArrayLike
+
+from amulet.registry import BiomePalette
+from .sub_chunk_array import SubChunkArrayContainer
+
+
+class Biome3DChunk:
+    def __init__(
+        self,
+        array_shape: tuple[int, int, int],
+        default_array: Union[int, ArrayLike],
+    ):
+        self.__biome_palette = BiomePalette()
+        self.__biomes = SubChunkArrayContainer(array_shape, default_array)
+
+    @property
+    def biome(self) -> SubChunkArrayContainer:
+        return self.__biomes
+
+    @biome.setter
+    def biome(
+        self,
+        sections: Iterable[int, ArrayLike],
+    ):
+        self.__biomes = SubChunkArrayContainer(
+            self.__biomes.array_shape, self.__biomes.default_array, sections
+        )
+
+    @property
+    def biome_palette(self):
+        return self.__biome_palette
+
+
+class Biome2DChunk:
+    def __init__(
+        self,
+        array_shape: tuple[int, int],
+        array: Union[int, ArrayLike],
+    ):
+        if (
+            not isinstance(array_shape, tuple)
+            and len(array_shape) == 2
+            and all(isinstance(s, int) for s in array_shape)
+        ):
+            raise TypeError
+
+        self.__array_shape = array_shape
+        self.__biome_palette = BiomePalette()
+        self.biome = array
+
+    @property
+    def biome(self) -> numpy.ndarray:
+        return self.__biomes
+
+    @biome.setter
+    def biome(
+        self,
+        array: Union[int, ArrayLike],
+    ):
+        if isinstance(array, int):
+            self.__biomes = numpy.full(self.__array_shape, array, dtype=numpy.uint32)
+        else:
+            array = numpy.asarray(array)
+            if not isinstance(array, numpy.ndarray):
+                raise TypeError
+            if array.shape != self.__array_shape or array.dtype != numpy.uint32:
+                raise ValueError
+            self.__biomes = numpy.array(array)
+
+    @property
+    def biome_palette(self):
+        return self.__biome_palette

amulet/chunk/block.py

@@ -0,0 +1,33 @@
+from typing import Union, Iterable
+
+from numpy.typing import ArrayLike
+
+from amulet.registry import BlockPalette
+from .sub_chunk_array import SubChunkArrayContainer
+
+
+class BlockChunk:
+    def __init__(
+        self,
+        array_shape: tuple[int, int, int],
+        default_array: Union[int, ArrayLike],
+    ):
+        self.__block_palette = BlockPalette()
+        self.__blocks = SubChunkArrayContainer(array_shape, default_array)
+
+    @property
+    def block(self) -> SubChunkArrayContainer:
+        return self.__blocks
+
+    @block.setter
+    def block(
+        self,
+        sections: Iterable[int, ArrayLike],
+    ):
+        self.__blocks = SubChunkArrayContainer(
+            self.__blocks.array_shape, self.__blocks.default_array, sections
+        )
+
+    @property
+    def block_palette(self):
+        return self.__block_palette

amulet/chunk/block_entity.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+from collections.abc import MutableMapping
+from typing import Iterable, Iterator
+
+from amulet.api.data_types import BlockCoordinates
+from amulet.block_entity import BlockEntity
+
+
+class BlockEntityContainer(MutableMapping[BlockCoordinates, BlockEntity]):
+    """
+    A MutableMapping that can only store :class:`BlockEntity` instances
+    under the absolute coordinate of the block entity ``Tuple[int, int, int]``
+    """
+
+    def __init__(
+        self,
+        block_entities: Iterable[BlockCoordinates, BlockEntity] = (),
+    ):
+        self._block_entities = {}
+
+        for location, block_entity in block_entities:
+            self[location] = block_entity
+
+    def __setitem__(self, coordinate: BlockCoordinates, block_entity: BlockEntity):
+        """
+        Set the :class:`BlockEntity` at ``coordinate``.
+
+        >>> block_entities: BlockEntityContainer
+        >>> x = y = z = 0
+        >>> block_entities[(x, y, z)] = block_entity
+
+        :param coordinate: The coordinate to set the block entity at.
+        :param block_entity: The block entity to set at the specified coordinate.
+        """
+        if (
+            not isinstance(coordinate, tuple)
+            and len(coordinate) == 3
+            and all(isinstance(c, int) for c in coordinate)
+        ):
+            raise TypeError
+        if not isinstance(block_entity, BlockEntity):
+            raise TypeError
+        self._block_entities[coordinate] = block_entity
+
+    def __delitem__(self, coordinate: BlockCoordinates):
+        """
+        Remove the :class:`BlockEntity` at ``coordinate``.
+
+        :param coordinate: The coordinate to remove the block entity from.
+        """
+        del self._block_entities[coordinate]
+
+    def __getitem__(self, coordinate: BlockCoordinates) -> BlockEntity:
+        """
+        Get the :class:`BlockEntity` at ``coordinate``.
+
+        >>> block_entities: BlockEntityContainer
+        >>> x = y = z = 0
+        >>> block_entity = block_entities[(x, y, z)]
+
+        :param coordinate: The coordinate to find the block entity at.
+        :return: The block entity at the specified coordinate.
+        :raises:
+            KeyError if there is no BlockEntity at the given coordinate.
+        """
+        return self._block_entities[coordinate]
+
+    def __len__(self) -> int:
+        return len(self._block_entities)
+
+    def __iter__(self) -> Iterator[BlockCoordinates]:
+        yield from self._block_entities
+
+
+class BlockEntityChunk:
+    """A chunk that supports block entities"""
+
+    def __init__(self):
+        self.__block_entity = BlockEntityContainer()
+
+    @property
+    def block_entity(self) -> BlockEntityContainer:
+        return self.__block_entity
+
+    @block_entity.setter
+    def block_entity(
+        self,
+        block_entities: Iterable[BlockCoordinates, BlockEntity],
+    ):
+        self.__block_entity = BlockEntityContainer(block_entities)

amulet/chunk/entity.py

@@ -0,0 +1,43 @@
+from typing import Iterable, Iterator
+from amulet.entity import Entity
+from collections.abc import MutableSet
+
+
+class EntityContainer(MutableSet[Entity]):
+    def __init__(self, entities: Iterable[Entity] = ()):
+        self._entities = set(entities)
+
+    def add(self, entity: Entity):
+        if not isinstance(entity, Entity):
+            raise TypeError("Expected an Entity")
+        self._entities.add(entity)
+
+    def discard(self, entity: Entity) -> None:
+        self._entities.discard(entity)
+
+    def __contains__(self, entity: Entity) -> bool:
+        return entity in self._entities
+
+    def __len__(self) -> int:
+        return len(self._entities)
+
+    def __iter__(self) -> Iterator[Entity]:
+        yield from self._entities
+
+
+class EntityChunk:
+    """A chunk that supports entities"""
+
+    def __init__(self):
+        self.__entity = EntityContainer()
+
+    @property
+    def entity(self) -> EntityContainer:
+        return self.__entity
+
+    @entity.setter
+    def entity(
+        self,
+        entities: Iterable[Entity],
+    ):
+        self.__entity = EntityContainer(entities)

amulet/chunk/sub_chunk_array.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+from collections.abc import MutableMapping
+from typing import Iterator, Union, Iterable
+
+import numpy
+from numpy.typing import ArrayLike
+
+
+class SubChunkArrayContainer(MutableMapping[int, numpy.ndarray]):
+    """A container of sub-chunk arrays"""
+
+    def __init__(
+        self,
+        array_shape: tuple[int, int, int],
+        default_array: Union[int, ArrayLike],
+        arrays: Iterable[int, ArrayLike] = (),
+    ):
+        self._shape = array_shape
+        self.default_array = default_array
+
+        self._arrays = {}
+        for cy, array in arrays:
+            self[cy] = array
+
+    @property
+    def array_shape(self) -> tuple[int, int, int]:
+        return self._shape
+
+    @property
+    def default_array(self) -> Union[int, numpy.ndarray]:
+        return self._default_array
+
+    @default_array.setter
+    def default_array(self, default_array: Union[int, ArrayLike]):
+        if isinstance(default_array, int):
+            self._default_array = default_array
+        else:
+            self._default_array = self._cast_array(default_array)
+
+    def _cast_array(self, array) -> numpy.ndarray:
+        array = numpy.asarray(array)
+        if not isinstance(array, numpy.ndarray):
+            raise TypeError
+        if array.shape != self._shape or array.dtype != numpy.uint32:
+            raise ValueError
+        return array
+
+    def populate(self, cy: int):
+        """Populate the section from the default array."""
+        if cy in self._arrays:
+            return
+        default_array = self._default_array
+        if isinstance(default_array, int):
+            arr = numpy.full(self.array_shape, default_array, dtype=numpy.uint32)
+        else:
+            arr = numpy.array(default_array, dtype=numpy.uint32)
+        self[cy] = arr
+
+    def __setitem__(self, cy: int, array: ArrayLike):
+        if not isinstance(cy, int):
+            raise TypeError
+        self._arrays[cy] = self._cast_array(array)
+
+    def __delitem__(self, cy: int) -> None:
+        del self._arrays[cy]
+
+    def __getitem__(self, cy: int) -> numpy.ndarray:
+        return self._arrays[cy]
+
+    def __len__(self) -> int:
+        return len(self._arrays)
+
+    def __iter__(self) -> Iterator[int]:
+        yield from self._arrays