PyNBT.py

import os
from pickle import unpack

class PyNBT:
    root = []
    TAG_END = 0
    TAG_BYTE = 1
    TAG_SHORT = 2
    TAG_INT = 3
    TAG_LONG = 4
    TAG_FLOAT = 5
    TAG_DOUBLE = 6
    TAG_BYTE_ARRAY = 7
    TAG_STRING = 8
    TAG_LIST = 9
    TAG_COMPOUND = 10
    
    def checkLength(self, string, expect):
        length = len(string)
        assert (length == expect), 'Expected ' + str(expect) + 'bytes, got ' + str(length)
        
    def readTriad(self, string: bytes) -> int:
        self.checkLength(string, 3)
        return unpack('>L', b'\x00' + string)[0]

    def writeTriad(self, value: int) -> bytes:
        return pack('>L', value)[1:]

    def readLTriad(self, string: bytes) -> int:
        self.checkLength(str, 3)
        return unpack('<L', b'\x00' + str)[0]

    def writeLTriad(self, value: int) -> bytes:
        return pack('<L', value)[0:-1]
    
    def readBool(self, b: bytes) -> int:
        return unpack('?', b)[0]

    def writeBool(self, b: int) -> bytes:
        return b'\x01' if b else b'\x00'
  
    def readByte(self, c: bytes) -> int:
        self.checkLength(c, 1)
        return unpack('>B', c)[0]
    
    def readSignedByte(self, c: bytes) -> int:
        self.checkLength(c, 1)
        return unpack('>b', c)[0]

    def writeByte(self, c: int) -> bytes:
        return pack(">B", c)
    
    def readShort(self, string: bytes) -> int:
        self.checkLength(string, 2)
        return unpack('>H', string)[0]

    def writeShort(self, value: int) -> bytes:
        return pack('>H', value)
    
    def readLShort(self, string: bytes) -> int:
        self.checkLength(string, 2)
        return unpack('<H', string)[0]

    def writeLShort(self, value: int) -> bytes:
        return pack('<H', value)
    
    def readInt(self, string: bytes) -> int:
        self.checkLength(string, 4)
        return unpack('>L', string)[0]

    def writeInt(self, value: int) -> bytes:
        return pack('>L', value)

    def readLInt(self, string: bytes) -> int:
        self.checkLength(string, 4)
        return unpack('<L', string)[0]

    def writeLInt(self, value: int) -> bytes:
        return pack('<L', value)

    def readFloat(self, string: bytes) -> int:
        self.checkLength(string, 4)
        return unpack('>f', string)[0]

    def writeFloat(self, value: int) -> bytes:
        return pack('>f', value)

    def readLFloat(self, string: bytes) -> int:
        self.checkLength(string, 4)
        return unpack('<f', string)[0]

    def writeLFloat(self, value: int) -> bytes:
        return pack('<f', value)

    def readDouble(self, string: bytes) -> int:
        self.checkLength(string, 8)
        return unpack('>d', string)[0]

    def writeDouble(self, value: int) -> bytes:
        return pack('>d', value)

    def readLDouble(self, string: bytes) -> int:
        self.checkLength(string, 8)
        return unpack('<d', string)[0]

    def writeLDouble(self, value: int) -> bytes:
        return pack('<d', value)

    def readLong(self, string: bytes) -> int:
        self.checkLength(string, 8)
        return unpack('>L', string)[0]

    def writeLong(self, value: int) -> bytes:
        return pack('>L', value)

    def readLLong(self, string: bytes) -> int:
        self.checkLength(string, 8)
        return unpack('<L', string)[0]

    def writeLLong(self, value: int) -> bytes:
        return pack('<L', value)
    
    def loadFile(self, filename):
        with open(filename, 'rb') as fp:
            bname = os.path.splitext(os.path.basename(filename))[0]
            if bname == 'level':
                version = self.readLInt(fp.read(4))
                lenght = self.readLInt(fp.read(4))
            elif bname == 'entities':
                fp.read(12)
            self.traverseTag(fp, self.root)
            return self.root[-1]
            
    def traverseTag(self, fp, tree):
        tagType = self.readType(fp, self.TAG_BYTE)
        if tagType == self.TAG_END:
            return False
        else:
            tagName = self.readType(fp, self.TAG_STRING)
            tagData = self.readType(fp, tagType)
            tree = []
            tree.append({'type': tagType, 'name': tagName, 'value': tagData})
            return True
        
    def readType(self, fp, tagType):
        if tagType == self.TAG_BYTE:
            return self.readByte(fp.read(1))
        elif tagType == self.TAG_SHORT:
            return self.readLShort(fp.read(2))
        elif tagType == self.TAG_INT:
            return self.readLInt(fp.read(4))
        elif tagType == self.TAG_LONG:
            return self.readLLong(fp.read(8))
        elif tagType == self.TAG_FLOAT:
            return self.readLFloat(fp.read(4))
        elif tagType == self.TAG_DOUBLE:
            return self.readLDouble(fp.read(8))
        elif tagType == self.TAG_BYTE_ARRAY:
            arrayLength = self.readType(fp, self.TAG_INT)
            arr = []
            i = 0
            while i < arrayLength:
                i += 1
                arr = self.readType(fp, self.TAG_BYTE)
                return arr
        elif tagType == self.TAG_STRING:
            stringLength = self.readType(fp, self.TAG_SHORT)
            if not stringLength:
                return ""
            string = fp.read(stringLength)
            return string
        elif tagType == self.TAG_LIST:
            tagID = self.readType(fp, self.TAG_BYTE)
            listLength = self.readType(fp, self.TAG_INT)
            list = {'type': tagID, 'value': []}
            i = 0
            while i < listLength:
                i += 1
                list["value"].append(self.readType(fp, tagID))
            return list
        elif tagType == self.TAG_COMPOUND:
            tree = []
            while self.traverseTag(fp, tree): pass
            return tree