Skip to content

Commit 4031db8

Browse files
gentlegiantJGCgentlegiantJGC
authored
Implement Sponge schematic v3 (#451)
* Add support for sponge schematic v3 * Add Schematic container tag * Fix blocks variable * Add entity Data tag * Fix reading block entities * Reformat * Enable configuration of schematic version
1 parent 3469b73 commit 4031db8

1 file changed

Lines changed: 122 additions & 57 deletions

File tree

‎amulet/level/formats/sponge_schem/format_wrapper.py‎

Lines changed: 122 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -47,13 +47,19 @@ class SpongeSchemWriteError(ObjectWriteError):
4747

4848
sponge_schem_interface = SpongeSchemInterface()
4949

50-
max_schem_version = 2
50+
max_schem_version = 3
5151

5252

5353
def _is_sponge(path: str):
5454
"""Check if a file is actually a sponge schematic file."""
5555
try:
56-
return "BlockData" in load_nbt(path).compound
56+
tag = load_nbt(path).compound
57+
schematic = tag.get("Schematic")
58+
if isinstance(schematic, CompoundTag):
59+
root = schematic
60+
else:
61+
root = tag
62+
return "Version" in root
5763
except:
5864
return False
5965

@@ -84,6 +90,7 @@ def _create(
8490
bounds: Union[
8591
SelectionGroup, Dict[Dimension, Optional[SelectionGroup]], None
8692
] = None,
93+
schematic_version: int = max_schem_version,
8794
**kwargs,
8895
):
8996
if not overwrite and os.path.isfile(self.path):
@@ -95,18 +102,24 @@ def _create(
95102
self._set_selection(bounds)
96103
self._is_open = True
97104
self._has_lock = True
105+
self._schem_version = schematic_version
98106

99107
def open_from(self, f: BinaryIO):
100-
sponge_schem = load_nbt(f).compound
108+
root_tag = load_nbt(f).compound
109+
schematic_tag = root_tag.get("Schematic")
110+
if isinstance(schematic_tag, CompoundTag):
111+
sponge_schem = schematic_tag
112+
else:
113+
sponge_schem = root_tag
101114
version_tag = sponge_schem.get("Version")
102115
if not isinstance(version_tag, IntTag):
103116
raise SpongeSchemReadError("Version key must exist and be an integer.")
104-
version = version_tag.py_int
117+
self._schem_version = version = version_tag.py_int
105118
if version == 1:
106119
raise SpongeSchemReadError(
107120
"Sponge Schematic Version 1 is not supported currently."
108121
)
109-
elif version == 2:
122+
elif 2 <= version <= 3:
110123
offset = sponge_schem.get("Offset")
111124
if isinstance(offset, IntArrayTag) and len(offset) == 3:
112125
min_point = numpy.array(offset)
@@ -138,12 +151,25 @@ def open_from(self, f: BinaryIO):
138151
self._platform = translator_version.platform
139152
self._version = translator_version.data_version
140153

141-
packed_block_data = sponge_schem.get("BlockData")
142-
if not isinstance(packed_block_data, ByteArrayTag):
154+
if version == 2:
155+
block_palette_tag = sponge_schem.get("Palette")
156+
block_index_tag = sponge_schem.get("BlockData")
157+
block_entities_tag = sponge_schem.get("BlockEntities")
158+
elif version == 3:
159+
blocks_tag = sponge_schem.get("Blocks")
160+
if not isinstance(blocks_tag, CompoundTag):
161+
raise SpongeSchemReadError("Blocks must be a CompoundTag")
162+
block_palette_tag = blocks_tag.get("Palette")
163+
block_index_tag = blocks_tag.get("Data")
164+
block_entities_tag = blocks_tag.get("BlockEntities")
165+
else:
166+
raise RuntimeError
167+
168+
if not isinstance(block_index_tag, ByteArrayTag):
143169
raise SpongeSchemReadError("BlockData must be a ByteArrayTag")
144170

145171
unpacked_block_data = decode_byte_array(
146-
numpy.array(packed_block_data, dtype=numpy.uint8)
172+
numpy.array(block_index_tag, dtype=numpy.uint8)
147173
)
148174
if len(unpacked_block_data) != numpy.prod(size):
149175
raise SpongeSchemReadError(
@@ -158,17 +184,11 @@ def open_from(self, f: BinaryIO):
158184
(2, 0, 1), # YZX => XYZ
159185
)
160186

161-
if "Palette" not in sponge_schem:
162-
raise SpongeSchemReadError(
163-
"Amulet is not able to read Sponge Schem files with no block palette."
164-
)
165-
166-
palette_data = sponge_schem.get("Palette")
167-
if not isinstance(palette_data, CompoundTag):
187+
if not isinstance(block_palette_tag, CompoundTag):
168188
raise SpongeSchemReadError("Palette must be a CompoundTag.")
169189

170190
block_palette: Dict[int, Block] = {}
171-
for blockstate, index_tag in palette_data.items():
191+
for blockstate, index_tag in block_palette_tag.items():
172192
index = index_tag.py_int
173193
if index in block_palette:
174194
raise SpongeSchemReadError(
@@ -207,29 +227,26 @@ def open_from(self, f: BinaryIO):
207227
[],
208228
)
209229

210-
if "BlockEntities" in sponge_schem:
211-
block_entities = sponge_schem["BlockEntities"]
212-
if not (
213-
isinstance(block_entities, ListTag)
214-
and (
215-
len(block_entities) == 0 or block_entities.list_data_type == 10
216-
) # CompoundTag.tag_id
217-
):
218-
raise SpongeSchemReadError(
219-
"BlockEntities must be a ListTag of compound tags."
220-
)
221-
222-
for block_entity in block_entities:
223-
if "Pos" not in block_entity:
224-
continue
225-
226-
pos_tag = block_entity["Pos"]
230+
if (
231+
isinstance(block_entities_tag, ListTag)
232+
and block_entities_tag.list_data_type == CompoundTag.tag_id
233+
):
234+
for block_entity in block_entities_tag:
235+
pos_tag = block_entity.get("Pos")
227236
if not (isinstance(pos_tag, IntArrayTag) and len(pos_tag) == 3):
228237
continue
229238

230239
pos = pos_tag.np_array + min_point
231240
x, y, z = pos
232241
block_entity["Pos"] = IntArrayTag(pos)
242+
if version == 2:
243+
extra = block_entity.pop("Extra", None)
244+
elif version == 3:
245+
extra = block_entity.pop("Data", None)
246+
else:
247+
raise RuntimeError
248+
if isinstance(extra, CompoundTag):
249+
block_entity.update(extra)
233250
cx, cz = x >> 4, z >> 4
234251
if (cx, cz) in self._chunks and (x, y, z) in self._chunks[
235252
(cx, cz)
@@ -249,10 +266,7 @@ def open_from(self, f: BinaryIO):
249266
)
250267

251268
for entity in entities:
252-
if "Pos" not in entity:
253-
continue
254-
255-
pos = entity["Pos"]
269+
pos = entity.get("Pos")
256270
if not (
257271
isinstance(pos, ListTag)
258272
and len(pos) == 3
@@ -272,6 +286,14 @@ def open_from(self, f: BinaryIO):
272286
IntTag(z),
273287
]
274288
)
289+
if version == 2:
290+
extra = entity.pop("Extra", None)
291+
elif version == 3:
292+
extra = entity.pop("Data", None)
293+
else:
294+
raise RuntimeError
295+
if isinstance(extra, CompoundTag):
296+
entity.update(extra)
275297
cx, cz = numpy.floor([x, z]).astype(int) >> 4
276298
if (cx, cz) in self._chunks and (x, y, z) in self._chunks[
277299
(cx, cz)
@@ -316,16 +338,16 @@ def save_to(self, f: BinaryIO):
316338
raise SpongeSchemReadError(
317339
"Sponge Schematic Version 1 is not supported currently."
318340
)
319-
elif self._schem_version == 2:
341+
elif 2 <= self._schem_version <= 3:
320342
selection = self._bounds[self.dimensions[0]].selection_boxes[0]
321343
if any(s > 2**16 - 1 for s in selection.shape):
322344
raise SpongeSchemWriteError(
323345
"The structure is too large to be exported to a Sponge Schematic file. It must be 2^16 - 1 at most in each dimension."
324346
)
325347
overflowed_shape = [s if s < 2**15 else s - 2**16 for s in selection.shape]
326-
tag = CompoundTag(
348+
schematic_tag = CompoundTag(
327349
{
328-
"Version": IntTag(2),
350+
"Version": IntTag(self._schem_version),
329351
"DataVersion": IntTag(self._version),
330352
"Width": ShortTag(overflowed_shape[0]),
331353
"Height": ShortTag(overflowed_shape[1]),
@@ -355,22 +377,42 @@ def save_to(self, f: BinaryIO):
355377
blocks[box.slice] = chunk.blocks + palette_len
356378
palette.append(chunk.palette)
357379
palette_len += len(chunk.palette)
358-
for be in chunk.block_entities:
359-
be = copy.deepcopy(be)
360-
be["Pos"] = IntArrayTag(be["Pos"].np_array - selection.min)
361-
block_entities.append(be)
362-
363-
for e in chunk.entities:
364-
e = copy.deepcopy(e)
365-
x, y, z = e["Pos"]
366-
e["Pos"] = ListTag(
380+
for block_entity in chunk.block_entities:
381+
block_entity = copy.deepcopy(block_entity)
382+
pos = IntArrayTag(
383+
block_entity.pop("Pos").np_array - selection.min
384+
)
385+
if self._schem_version == 2:
386+
block_entity["Pos"] = pos
387+
elif self._schem_version == 3:
388+
id_ = block_entity.pop("Id")
389+
block_entity = CompoundTag(
390+
{"Pos": pos, "Id": id_, "Data": block_entity}
391+
)
392+
else:
393+
raise RuntimeError
394+
block_entities.append(block_entity)
395+
396+
for entity in chunk.entities:
397+
entity = copy.deepcopy(entity)
398+
x, y, z = entity["Pos"]
399+
pos = ListTag(
367400
[
368401
IntTag(x - selection.min_x),
369402
IntTag(y - selection.min_y),
370403
IntTag(z - selection.min_z),
371404
]
372405
)
373-
entities.append(e)
406+
if self._schem_version == 2:
407+
entity["Pos"] = pos
408+
elif self._schem_version == 3:
409+
id_ = entity.pop("Id")
410+
entity = CompoundTag(
411+
{"Pos": pos, "Id": id_, "Data": entity}
412+
)
413+
else:
414+
raise RuntimeError
415+
entities.append(entity)
374416

375417
compact_palette, lut = brute_sort_objects_no_hash(
376418
numpy.concatenate(palette)
@@ -381,20 +423,43 @@ def save_to(self, f: BinaryIO):
381423
block: Block
382424
block_palette.append(block.blockstate)
383425

384-
tag["PaletteMax"] = IntTag(len(compact_palette))
385-
tag["Palette"] = CompoundTag(
426+
block_palette_tag = CompoundTag(
386427
{
387428
blockstate: IntTag(index)
388429
for index, blockstate in enumerate(block_palette)
389430
}
390431
)
391-
tag["BlockData"] = ByteArrayTag(list(encode_array(blocks)))
392-
if block_entities:
393-
tag["BlockEntities"] = ListTag(block_entities)
432+
block_index_tag = ByteArrayTag(list(encode_array(blocks)))
433+
434+
if self._schem_version == 2:
435+
schematic_tag["PaletteMax"] = IntTag(len(compact_palette))
436+
schematic_tag["Palette"] = block_palette_tag
437+
schematic_tag["BlockData"] = block_index_tag
438+
if block_entities:
439+
schematic_tag["BlockEntities"] = ListTag(block_entities)
440+
elif self._schem_version == 3:
441+
blocks_tag = CompoundTag(
442+
{
443+
"Palette": block_palette_tag,
444+
"Data": block_index_tag,
445+
}
446+
)
447+
if block_entities:
448+
blocks_tag["BlockEntities"] = ListTag(block_entities)
449+
schematic_tag["Blocks"] = blocks_tag
450+
else:
451+
raise RuntimeError
452+
394453
if entities:
395-
tag["Entities"] = ListTag(entities)
454+
schematic_tag["Entities"] = ListTag(entities)
396455

397-
NamedTag(tag, "Schematic").save_to(f)
456+
if self._schem_version == 2:
457+
named_tag = NamedTag(schematic_tag, "Schematic")
458+
elif self._schem_version == 3:
459+
named_tag = NamedTag(CompoundTag({"Schematic": schematic_tag}), "")
460+
else:
461+
raise RuntimeError
462+
named_tag.save_to(f)
398463
else:
399464
raise SpongeSchemReadError(
400465
f"Sponge Schematic Version {self._schem_version} is not supported currently."

0 commit comments

Comments
 (0)