Add an option to control which chunks will be modified (#453)

Author: gentlegiantJGC

amulet/api/level/base_level/base_level.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import time
 from typing import Union, Generator, Optional, Tuple, Callable, Set, Iterable
 import traceback
 import numpy
@@ -32,7 +31,7 @@
 from amulet.utils.world_utils import block_coords_to_chunk_coords
 from .chunk_manager import ChunkManager
 from amulet.api.history.history_manager import MetaHistoryManager
-from .clone import clone
+from .clone import clone, PasteRule
 from amulet.api import wrapper as api_wrapper, level as api_level
 import PyMCTranslate
 from amulet.api.player import Player
@@ -657,6 +656,7 @@ def paste(
         include_entities: bool = True,
         skip_blocks: Tuple[Block, ...] = (),
         copy_chunk_not_exist: bool = False,
+        paste_rule: PasteRule = PasteRule.PasteAll,
     ):
         """Paste a level into this level at the given location.
         Note this command may change in the future.
@@ -672,6 +672,7 @@ def paste(
         :param include_entities: Include entities when pasting the structure.
         :param skip_blocks: If a block matches a block in this list it will not be copied.
         :param copy_chunk_not_exist: If a chunk does not exist in the source should it be copied over as air. Always False where level is a World.
+        :param paste_rule: Control which chunks can be pasted into.
         :return:
         """
         return generator_unpacker(
@@ -687,6 +688,7 @@ def paste(
                 include_entities,
                 skip_blocks,
                 copy_chunk_not_exist,
+                paste_rule,
             )
         )
 
@@ -703,6 +705,7 @@ def paste_iter(
         include_entities: bool = True,
         skip_blocks: Tuple[Block, ...] = (),
         copy_chunk_not_exist: bool = False,
+        paste_rule: PasteRule = PasteRule.PasteAll,
     ) -> Generator[float, None, None]:
         """Paste a structure into this structure at the given location.
         Note this command may change in the future.
@@ -718,6 +721,7 @@ def paste_iter(
         :param include_entities: Include entities when pasting the structure.
         :param skip_blocks: If a block matches a block in this list it will not be copied.
         :param copy_chunk_not_exist: If a chunk does not exist in the source should it be copied over as air. Always False where level is a World.
+        :param paste_rule: Control which chunks can be pasted into.
         :return: A generator of floats from 0 to 1 with the progress of the paste operation.
         """
         yield from clone(
@@ -734,6 +738,7 @@ def paste_iter(
             include_entities,
             skip_blocks,
             copy_chunk_not_exist,
+            paste_rule,
         )
 
     def get_version_block(

amulet/api/level/base_level/clone.py

@@ -2,6 +2,7 @@
 from typing import TYPE_CHECKING, Tuple, Generator, Optional
 import numpy
 import os
+from enum import Enum
 
 from amulet.api.data_types import Dimension, BlockCoordinates, FloatTriplet
 from amulet.api.selection import SelectionGroup, SelectionBox
@@ -42,6 +43,12 @@ def gen_paste_blocks(
     )
 
 
+class PasteRule(Enum):
+    PasteAll = "PasteAll"  # Paste in all chunk.
+    PasteExist = "PasteExist"  # Paste only in chunks that exist.
+    PasteNotExist = "PasteNotExist"  # Paste only in chunks that do not exist.
+
+
 def clone(
     src_structure: "BaseLevel",
     src_dimension: Dimension,
@@ -56,6 +63,7 @@ def clone(
     include_entities: bool = True,
     skip_blocks: Tuple[Block, ...] = (),
     copy_chunk_not_exist: bool = False,
+    paste_rule: PasteRule = PasteRule.PasteAll,
 ) -> Generator[float, None, None]:
     """Clone the source object data into the destination object with an optional transform.
     The src and dst can be the same object.
@@ -73,6 +81,7 @@ def clone(
     :param include_entities: Include entities from the `src_structure`.
     :param skip_blocks: If a block matches a block in this list it will not be copied.
     :param copy_chunk_not_exist: If a chunk does not exist in the source should it be copied over as air. Always False where `src_structure` is a World.
+    :param paste_rule: Control which chunks can be pasted into.
     :return: A generator of floats from 0 to 1 with the progress of the paste operation.
     """
     location = tuple(location)
@@ -106,6 +115,10 @@ def clone(
 
         src_structure: "BaseLevel"
 
+        # If the paste mode is PasteNotExist we need to create the chunk and track that we created it
+        # so that future modifications in this operation are allowed.
+        created_chunks = set[tuple[int, int]]()
+
         # TODO: I don't know if this is feasible for large boxes: get the intersection of the source and destination selections and iterate over that to minimise work
         if any(rotation) or any(s != 1 for s in scale):
             # if the selection needs transforming
@@ -137,6 +150,7 @@ def clone(
                     for progress, src_coords, dst_coords in box.transformed_points(
                         transform
                     ):
+                        yield sum_progress + volumes[box_index] * progress
                         if src_coords is not None:
                             dst_cx, dst_cy, dst_cz = dst_coords[0] >> 4
                             if (dst_cx, dst_cz) != last_dst:
@@ -146,11 +160,20 @@ def clone(
                                         dst_cx, dst_cz, dst_dimension
                                     )
                                 except ChunkDoesNotExist:
+                                    if paste_rule == PasteRule.PasteExist:
+                                        continue
                                     dst_chunk = dst_structure.create_chunk(
                                         dst_cx, dst_cz, dst_dimension
                                     )
+                                    created_chunks.add((dst_cx, dst_cz))
                                 except ChunkLoadError:
                                     dst_chunk = None
+                                else:
+                                    if (
+                                        paste_rule == PasteRule.PasteNotExist
+                                        and (dst_cx, dst_cz) not in created_chunks
+                                    ):
+                                        continue
 
                             src_coords = numpy.floor(src_coords).astype(int)
                             # due to how the coords are found dst_coords will all be in the same sub-chunk
@@ -252,7 +275,6 @@ def clone(
                                             if location in dst_chunk.block_entities:
                                                 del dst_chunk.block_entities[location]
                                         dst_chunk.changed = True
-                        yield sum_progress + volumes[box_index] * progress
                     sum_progress += volumes[box_index]
 
         else:
@@ -297,16 +319,27 @@ def clone(
                 dst_slices: Tuple[slice, slice, slice]
                 dst_box: SelectionBox
 
+                yield count / iter_count
+                count += 1
+
                 # load the destination chunk
                 try:
                     dst_chunk = dst_structure.get_chunk(dst_cx, dst_cz, dst_dimension)
                 except ChunkDoesNotExist:
+                    if paste_rule == PasteRule.PasteExist:
+                        continue
                     dst_chunk = dst_structure.create_chunk(
                         dst_cx, dst_cz, dst_dimension
                     )
+                    created_chunks.add((dst_cx, dst_cz))
                 except ChunkLoadError:
-                    count += 1
                     continue
+                else:
+                    if (
+                        paste_rule == PasteRule.PasteNotExist
+                        and (dst_cx, dst_cz) not in created_chunks
+                    ):
+                        continue
 
                 if include_blocks:
                     # a boolean array specifying if each index should be pasted.
@@ -383,7 +416,4 @@ def clone(
                     # TODO: implement pasting entities when we support entities
                     pass
 
-                count += 1
-                yield count / iter_count
-
         yield 1.0

amulet/operations/paste.py

@@ -4,6 +4,7 @@
 
 from amulet.api.data_types import Dimension, BlockCoordinates, FloatTriplet
 from amulet.api.block import Block, UniversalAirLikeBlocks
+from amulet.api.level.base_level.clone import PasteRule
 
 if TYPE_CHECKING:
     from amulet.api.level import BaseLevel
@@ -20,6 +21,7 @@ def paste(
     copy_air=True,
     copy_water=True,
     copy_lava=True,
+    paste_rule: PasteRule = PasteRule.PasteAll,
 ):
     for _ in paste_iter(
         dst,
@@ -32,6 +34,7 @@ def paste(
         copy_air,
         copy_water,
         copy_lava,
+        paste_rule,
     ):
         pass
 
@@ -47,6 +50,7 @@ def paste_iter(
     copy_air=True,
     copy_water=True,
     copy_lava=True,
+    paste_rule: PasteRule = PasteRule.PasteAll,
 ):
     yield from dst.paste_iter(
         src,
@@ -62,4 +66,5 @@ def paste_iter(
         + (Block("universal_minecraft", "water"),) * bool(not copy_water)
         + (Block("universal_minecraft", "lava"),) * bool(not copy_lava),
         True,
+        paste_rule,
     )