Support parsing block_states with version greater than 1.18

mcworldlib/chunk.py

@@ -44,8 +44,36 @@ def entities(self):
     def entities(self, value: nbt.List[nbt.Compound]):
         self.data_root['Entities'] = value
 
-    def get_blocks(self):
-        """Yield a (Y, Palette, BlockState Indexes Array) tuple for every chunk section.
+
+    def is_version_1_21(self):
+        if "sections" in self.data_root:
+            return True
+        else:
+            return False
+
+
+    def get_blocks_1_21(self):
+        """Yield a (Y, palette, block_states Indexes Array) tuple for every chunk section.
+
+        Y: Y "level" of the section, the section "index" (NOT the NBT section index!)
+        palette, Indexes: See get_section_blocks()
+        """
+        blocks = {}
+        for section in self.data_root['sections']:
+            # noinspection PyPep8Naming
+            Y = int(section['Y'])
+            palette, indexes = self.get_section_blocks(Y, _section=section)
+            if palette:
+                blocks[Y] = palette, indexes
+        for Y in sorted(blocks):
+            # noinspection PyRedundantParentheses
+            yield (Y, *blocks[Y])
+
+
+    def get_blocks_old(self):
+        """For version lower than 1.21.1
+
+        Yield a (Y, Palette, BlockState Indexes Array) tuple for every chunk section.
 
         Y: Y "level" of the section, the section "index" (NOT the NBT section index!)
         Palette, Indexes: See get_section_blocks()
@@ -61,9 +89,54 @@ def get_blocks(self):
             # noinspection PyRedundantParentheses
             yield (Y, *blocks[Y])
 
-    # noinspection PyPep8Naming
+
+    def get_blocks(self):
+        if self.is_version_1_21():
+            return self.get_blocks_1_21()
+        else:
+            return self.get_blocks_old()
+
+
     def get_section_blocks(self, Y: int, _section=None):
-        """Return a (Palette, BlockState Indexes Array) tuple for a chunk section.
+        if self.is_version_1_21():
+            return self.get_section_blocks_1_21(Y, _section)
+        else:
+            return self.get_section_blocks_old(Y, _section)
+
+
+    # noinspection PyPep8Naming
+    def get_section_blocks_1_21(self, Y: int, _section=None):
+        """Return a (palette, block_state Indexes Array) tuple for a chunk section.
+
+        palette: NBT List of Block States, straight from NBT data
+        Indexes: 16 x 16 x 16 numpy.ndarray, in YZX order, of indexes matching palette's
+        """
+        section = _section
+        if not section:
+            for section in self.data_root.get('sections', []):
+                if section.get('Y') == Y:
+                    break
+            else:
+                return None, None
+
+        if 'block_states' not in section:
+            return None, None
+
+        states = section['block_states']
+
+        palette = states['palette']
+
+        if 'data' not in states:
+            return palette, numpy.zeros((u.SECTION_HEIGHT, *reversed(u.CHUNK_SIZE)))
+
+        indexes = self._decode_blockstates(states['data'], palette)
+        return palette, indexes.reshape((u.SECTION_HEIGHT, *reversed(u.CHUNK_SIZE)))
+
+
+    def get_section_blocks_old(self, Y: int, _section=None):
+        """For version lower than 1.21.1
+
+        Return a (Palette, BlockState Indexes Array) tuple for a chunk section.
 
         Palette: NBT List of Block States, straight from NBT data
         Indexes: 16 x 16 x 16 numpy.ndarray, in YZX order, of indexes matching Palette's
@@ -83,6 +156,7 @@ def get_section_blocks(self, Y: int, _section=None):
         indexes = self._decode_blockstates(section['BlockStates'], palette)
         return palette, indexes.reshape((u.SECTION_HEIGHT, *reversed(u.CHUNK_SIZE)))
 
+
     def _decode_blockstates(self, data, palette=None):
         """Decode an NBT BlockStates LongArray to a block state indexes array"""
         pack_bits = data.itemsize * 8  # 64 bits for each Long Array element
@@ -99,17 +173,18 @@ def bits_from_data(): return len(data) * pack_bits // self.BS_INDEXES
             return bit_length
 
         bits = bits_per_index()
+        padding_per_long = 64 % bits
         # Adapted from Amulet-Core's decode_long_array()
         # https://github.com/Amulet-Team/Amulet-Core/blob/develop/amulet/utils/world_utils.py
         indexes = numpy.packbits(
             numpy.pad(
                 numpy.unpackbits(
-                        data[::-1].astype(f">i{pack_bits//8}").view(f"uint{pack_bits//8}")
-                    ).reshape(-1, bits),
+                        data[::-1].astype(f">i{pack_bits//8}").view(f"uint{pack_bits//8}").unpack()
+                    ).reshape(-1, 64)[:, padding_per_long:64].reshape(-1, bits),
                 [(0, 0), (pack_bits - bits, 0)],
                 "constant",
             )
-        ).view(dtype=">q")[::-1]
+        ).view(dtype=">q")[::-1][:4096]
         return indexes
 
     # noinspection PyMethodMayBeStatic

mcworldlib/tree.py

@@ -209,9 +209,9 @@ def is_nbt_container(tag: 'nbt.AnyTag') -> bool:
 def tests():
     import json
     for data in (
-        json.load(open("../data/New World/advancements/"
+        json.load(open("data/New World/advancements/"
                        "8b4accb8-d952-4050-97f2-e00c4423ba92.json")),
-        nbt.load_dat("../data/New World/level.dat"),
+        nbt.load_dat("data/New World/level.dat"),
         [{"x": 1, "y": 2}, "a", ((4, {"z": 5}, "b"), 6, "c")],
         {"name": ["Rodrigo", ("E", "S"), "Silva"], "alias": "MestreLion"}
     ):

tests/tests.py

@@ -1,13 +1,16 @@
 #!/usr/bin/env python3
+"""
+Run tests by `python tests/tests.py`
+"""
 
 import io
 import logging
 import os.path
 import sys
 
+sys.path.append(os.path.abspath("."))
 import mcworldlib as mc
 
-
 logging.basicConfig(level=logging.DEBUG, format='%(levelname)-6s: %(message)s')
 
 
@@ -80,7 +83,7 @@ def diamonds():
 
 
 def new_diamonds():
-    world = mc.load('New World')
+    world = mc.load('data/New World')
 
     for item in world.level['Data']['Player']['Inventory']:
         item['id'] = mc.String('minecraft:diamond_block')
@@ -128,7 +131,7 @@ def player_pos():
 
 def blocks(w=None):
     if w is None:
-        w = mc.load('New World')
+        w = mc.load('data/New World')
 
     for chunk in [w.player.get_chunk()]:  # w.get_chunks(False):
         print(chunk)