minecraft_world_edit.py

import io
import nbt  # sudo python3 -m pip install NBT
import os


class BitReader:
    def __init__(self, data, size):
        self.data = data
        self.size = size
        self.total = len(self.data) * self.size
        self.index = 0

    def get_bit(self):
        assert self.index < self.total
        div, mod = divmod(self.index, self.size)
        # print(self.index, self.size, div, mod)
        result = (self.data[div] >> (self.size - 1 - mod)) & 1
        self.index += 1
        return result

    def get_bits(self, size):
        result = 0
        for _ in range(size):
            result = result << 1 | self.get_bit()
        return result


class BitWriter:
    def __init__(self, size):
        self.data = []
        self.size = size
        self.index = 0

    def write_bit(self, value):
        if self.index == len(self.data) * self.size:
            self.data.append(0)
        div, mod = divmod(self.index, self.size)
        bit_pos = self.size - 1 - mod
        self.data[div] = self.data[div] | (value << bit_pos)
        self.index += 1

    def write_int(self, value, size):
        for index in range(size):
            bit_pos = size - 1 - index
            self.write_bit((value >> bit_pos) & 1)

    def compliment(self, value):
        return value-(1<<(value.bit_length()))

    def get(self):
        return [
            value if value < 2 ** 63 else self.compliment(value)
            for value in self.data
        ]



class Chunk:
    @staticmethod
    def get_palette_index_size(count):
        # for index in [3, 8, 15, 16, 24, 31, 32, 33]: print(index, get_palette_index_size(index))
        size = 4
        while 2 ** size - 1 < count:
            size += 1
        return size

    @classmethod
    def load(cls, nbt_chunk):
        layers = {}
        for section in nbt_chunk["Level"]["Sections"]:
            palette = {}
            if "Palette" not in section:
                continue
            for index, item in enumerate(section["Palette"]):
                palette[index] = item["Name"].value
            palette_size = len(section["Palette"])
            palette_index_size = cls.get_palette_index_size(palette_size)
            bit_reader = BitReader(section["BlockStates"].value, size=64)
            yy_offset = section["Y"].value * 16
            for yy in range(16):
                layer = {}
                for xx in range(16):
                    for zz in range(16):
                        palette_index = bit_reader.get_bits(palette_index_size)
                        layer[(xx, zz)] = palette[palette_index]
                layers[yy_offset + yy] = layer
        return cls(nbt_chunk, layers)

    def __init__(self, nbt_chunk, layers):
        # print(nbt_chunk.pretty_tree())
        self._nbt_chunk = nbt_chunk
        self._layers = layers

    def serialize(self):
        air_block = "minecraft:air"
        buffer = io.BytesIO()
        sections = nbt.nbt.TAG_List(type=nbt.nbt.TAG_Compound)
        for subchunk_y in range(16):
            section = nbt.nbt.TAG_Compound()
            section["Y"] = nbt.nbt.TAG_Short(subchunk_y)
            palette = {}
            for offset_y in range(16):
                yy = subchunk_y * 16 + offset_y
                if yy not in self._layers:
                    if air_block not in palette:
                        palette[air_block] = len(palette)
                    continue
                for xx in range(16):
                    for zz in range(16):
                        value = self._layers[yy][(xx, zz)]
                        if value not in palette:
                            palette[value] = len(palette)
            section["Palette"] = nbt.nbt.TAG_List(type=nbt.nbt.TAG_Compound)
            for value, index in sorted(palette.items(), key=lambda pair: pair[1]):
                block = nbt.nbt.TAG_Compound()
                block["Name"] = nbt.nbt.TAG_String(value, "Name")
                section["Palette"].insert(index, block)
            palette_index_size = self.get_palette_index_size(len(palette))
            bit_writer = BitWriter(size=64)
            for offset_y in range(16):
                yy = subchunk_y * 16 + offset_y
                if yy in self._layers:
                    for xx in range(16):
                        for zz in range(16):
                            palette_index = palette[self._layers[yy][(xx, zz)]]
                            bit_writer.write_int(palette_index, palette_index_size)
                else:
                    for xx in range(16):
                        for zz in range(16):
                            palette_index = palette[air_block]
                            bit_writer.write_int(palette_index, palette_index_size)

            block_states = nbt.nbt.TAG_Long_Array()
            block_states.value = bit_writer.get()
            section["BlockStates"] = block_states
            sections.insert(subchunk_y, section)
        self._nbt_chunk["Sections"] = sections
        # return self._nbt_chunk["Sections"].pretty_tree()

        self._nbt_chunk.write_file(buffer=buffer)
        return buffer.getvalue()

    def get_block(self, xx, yy, zz):
        return self._layers[yy][(xx, zz)]

    def set_block(self, xx, yy, zz, value):
        self._layers[yy][(xx, zz)] = value

    def __str__(self):
        result = ""
        width = 10
        for yy in range(8):
            result += f"Layer {yy}\n"
            if yy not in self._layers:
                result += "Skipped\n\n"
                continue
            for xx in range(16):
                for zz in range(16):
                    name = self._layers[yy][(xx, zz)][10:]  # remove "minecraft:"
                    result += name[:width-1].ljust(width, " ")
                result += "\n"
            result += "\n"
        return result



class World:
    @classmethod
    def load(cls, path):
        nbt_world = nbt.world.AnvilWorldFolder(path)
        return cls(nbt_world)

    def __init__(self, nbt_world):
        self._nbt_world = nbt_world
        self._chunks = {}

    def get_chunk(self, cx, cz):
        key = (cx, cz)
        if key not in self._chunks:
            self._chunks[key] = Chunk.load(self._nbt_world.get_nbt(cx, cz))
        return self._chunks[key]

    def get_block(self, x, y, z):
        xd, xm = divmod(x, 16)
        zd, zm = divmod(z, 16)
        self.get_chunk(xd, zd)






prefix = "/Users/kaustubh/Library/Application Support/minecraft/saves/"
save_name = "Hacky Creative"
location = os.path.join(prefix, save_name)

if True:
    world = World.load(location)
    chunk = world.get_chunk(0, 0)
    # chunk.set_block(0, 0, 0, "minecraft:stone")
    nbt_chunk = nbt.nbt.NBTFile(buffer=io.BytesIO(chunk.serialize()))
    print(Chunk.load(nbt_chunk))
else:
    world = nbt.world.AnvilWorldFolder(os.path.join(prefix, save_name))
    chunk_nbt = world.get_nbt(0, 0)
    print(chunk_nbt.pretty_tree())