Permalink
Find file
Fetching contributors…
Cannot retrieve contributors at this time
635 lines (513 sloc) 17.4 KB
# cython: profile=True
# vim:set sw=2 sts=2 ts=2:
"""
Cython implementation
Named Binary Tag library. Serializes and deserializes TAG_* objects
to and from binary data. Load a Minecraft level by calling nbt.load().
Create your own TAG_* objects and set their values.
Save a TAG_* object to a file or StringIO object.
Read the test functions at the end of the file to get started.
This library requires Numpy. Get it here:
http://new.scipy.org/download.html
Official NBT documentation is here:
http://www.minecraft.net/docs/NBT.txt
Copyright 2010 David Rio Vierra
"""
import collections
import itertools
import struct
import gzip
from cStringIO import StringIO
from cpython cimport PyTypeObject, PyObject_TypeCheck, PyUnicode_DecodeUTF8, PyList_Append
cdef extern from "cStringIO.h":
struct PycStringIO_CAPI:
int cwrite (object o, char *buf, Py_ssize_t len)
PyTypeObject * OutputType
cdef extern from "cobject.h":
void * PyCObject_Import(char *module_name, char *cobject_name)
cdef PycStringIO_CAPI *PycStringIO = <PycStringIO_CAPI *>PyCObject_Import("cStringIO", "cStringIO_CAPI")
cdef PyTypeObject * StringO = PycStringIO.OutputType
cdef cwrite(obj, char * buf, size_t len):
#print "cwrite %s %s %d" % (map(ord, buf[:min(4, len)]), buf[:min(4, len)].decode('ascii', 'replace'), len)
return PycStringIO.cwrite(obj, buf, len)
import sys
import os
from os.path import exists
from contextlib import closing
from numpy import array, zeros, uint8, fromstring, ndarray, frombuffer
cimport numpy as np
cdef char TAG_END = 0
cdef char TAG_BYTE = 1
cdef char TAG_SHORT = 2
cdef char TAG_INT = 3
cdef char TAG_LONG = 4
cdef char TAG_FLOAT = 5
cdef char TAG_DOUBLE = 6
cdef char TAG_BYTE_ARRAY = 7
cdef char TAG_STRING = 8
cdef char TAG_LIST = 9
cdef char TAG_COMPOUND = 10
cdef char TAG_INT_ARRAY = 11
cdef char TAG_SHORT_ARRAY = 12
class NBTFormatError (ValueError):
pass
cdef class TAG_Value:
cdef unicode _name
cdef public char tagID
def __str__(self):
return self.tostr()
cdef tostr(self):
return str(self.__class__) + ": " + str(self.value)
property name:
def __get__(self):
return self._name
def __set__(self, val):
if isinstance(val, str): val = PyUnicode_DecodeUTF8(val, len(val), "strict")
self._name = val
def __reduce__(self):
return self.__class__, (self.value, self._name)
cdef class TAG_Number(TAG_Value):
pass
cdef class TAG_Array(TAG_Value):
pass
cdef class TAG_Byte(TAG_Number):
cdef public char value
cdef save_value(self, buf):
save_byte(self.value, buf)
def __init__(self, char value=0, name = u""):
self.value = value
self.name = name
self.tagID = TAG_BYTE
cdef class TAG_Short(TAG_Number):
cdef public short value
cdef save_value(self, buf):
save_short(self.value, buf)
def __init__(self, short value=0, name = u""):
self.value = value
self.name = name
self.tagID = TAG_SHORT
cdef class TAG_Int(TAG_Number):
cdef public int value
cdef save_value(self, buf):
save_int(self.value, buf)
def __init__(self, int value=0, name = u""):
self.value = value
self.name = name
self.tagID = TAG_INT
cdef class TAG_Long(TAG_Number):
cdef public long long value
cdef save_value(self, buf):
save_long(self.value, buf)
def __init__(self, long long value=0, name = u""):
self.value = value
self.name = name
self.tagID = TAG_LONG
cdef class TAG_Float(TAG_Number):
cdef public float value
cdef save_value(self, buf):
save_float(self.value, buf)
def __init__(self, float value=0., name = u""):
self.value = value
self.name = name
self.tagID = TAG_FLOAT
cdef class TAG_Double(TAG_Number):
cdef public double value
cdef save_value(self, buf):
save_double(self.value, buf)
def __init__(self, double value=0., name = u""):
self.value = value
self.name = name
self.tagID = TAG_DOUBLE
cdef class TAG_Byte_Array(TAG_Array):
cdef public object value
def __init__(self, value = zeros((0,), 'uint8'), name = u""):
self.value = value
self.name = name
self.tagID = TAG_BYTE_ARRAY
cdef save_value(self, buf):
save_byte_array(self.value, buf)
cdef class TAG_String(TAG_Value):
cdef unicode _value
def __init__(self, value = u"", name = u""):
if isinstance(value, str): value = PyUnicode_DecodeUTF8(value, len(value), "strict")
self.value = value
self.name = name
self.tagID = TAG_STRING
property value:
def __get__(self):
return self._value
def __set__(self, value):
if isinstance(value, str): value = PyUnicode_DecodeUTF8(value, len(value), "strict")
self._value = value
cdef save_value(self, buf):
save_string(self.value.encode('utf-8'), buf)
cdef class _TAG_List(TAG_Value):
cdef public list value
def __init__(self, value = None, name = u""):
self.value = list(value or [])
self.name = name
self.tagID = TAG_LIST
"""collection methods"""
def __getitem__(self, key):
return self.value[key]
def __setitem__(self, key, val):
self.value[key] = val
def __iter__(self):
return iter(self.value)
def __len__(self): return len(self.value)
def insert(self, idx, key):
self.value.insert(idx, key)
property list_type:
def __get__(self):
if len(self.value): return self.value[0].tagID
return TAG_BYTE
cdef save_value(self, buf):
save_tag_id(self.list_type, buf)
save_int(len(self.value), buf)
items = self.value
for subtag in items:
if subtag.tagID != self.list_type:
raise NBTFormatError, "Asked to save TAG_List with different types! Found %s and %s" % (subtag.tagID, self.list_type)
save_tag_value(subtag, buf)
class TAG_List(_TAG_List, collections.MutableSequence):
pass
cdef class _TAG_Compound(TAG_Value):
cdef public dict value
def __init__(self, value = None, name = u""):
self.value = value or {}
self.name = name
self.tagID = TAG_COMPOUND
"""collection methods"""
def __getitem__(self, key):
return self.value[key]
def __setitem__(self, key, val):
assert isinstance(val, TAG_Value)
val.name = key
self.value[key] = val
def __delitem__(self, key):
del self.value[key]
def __iter__(self): return iter(self.value)
def __contains__(self, k):return k in self.value
def __len__(self): return len(self.value)
def __str__(self):
return str(self.__class__) + ": " + str(self.value)
__repr__ = __str__
def add(self, tag):
assert tag.name
self[tag.name] = tag
cdef save_value(self, buf):
i = self.iteritems()
for name, subtag in i:
#print "save_tag_name", name, subtag.tagID, "Named", subtag.name,
save_tag_id(subtag.tagID, buf)
#print "id",
save_tag_name(subtag, buf)
#print "name",
save_tag_value(subtag, buf)
#print "value", name
save_tag_id(TAG_END, buf)
class TAG_Compound(_TAG_Compound, collections.MutableMapping):
def __init__(self, value = None, name = u""):
_TAG_Compound.__init__(self, value, name)
def save(self, filename = "", buf = None):
save_root_tag(self, filename, buf)
def saveGzipped(self, filename, compresslevel=1):
save_root_tag(self, filename)
cdef class TAG_Int_Array(TAG_Array):
cdef char _tagID(self): return TAG_INT_ARRAY
cdef public object value
cdef save_value(self, buf):
pass
cdef class TAG_Short_Array(TAG_Array):
cdef char _tagID(self): return TAG_SHORT_ARRAY
cdef public object value
cdef save_value(self, buf):
pass
#cdef int needswap = (sys.byteorder == "little")
cdef swab(void * vbuf, int nbytes):
cdef unsigned char * buf = <unsigned char *> vbuf
#print "Swapping ", nbytes, "bytes"
#for i in range(nbytes): print buf[i],
#print "to",
#if not needswap: return
cdef int i
for i in range((nbytes+1)/2):
buf[i], buf[nbytes-i-1] = buf[nbytes-i-1], buf[i]
#for i in range(nbytes): print buf[i],
import zlib
def gunzip(data):
#strip off the header and use negative WBITS to tell zlib there's no header
return zlib.decompress(data[10:], -zlib.MAX_WBITS)
def try_gunzip(data):
try:
data = gunzip(data)
except Exception, e:
pass
return data
def load(buf=None, filename=None):
try:
if isinstance(buf, basestring) and exists(buf):
filename = buf
except TypeError:
pass
if filename and exists(filename):
data = file(filename, "rb").read()
data = try_gunzip(data)
return load_buffer(data)
return load_buffer(try_gunzip(buf))
cdef class load_ctx:
cdef size_t offset
cdef char * buffer
cdef size_t size
cdef int require(self, size_t s) except -1:
#print "Asked for ", s
if s > self.size - self.offset:
raise NBTFormatError, "NBT Stream too short. Asked for %d, only had %d" % (s, (self.size - self.offset))
return 0
should_dump = False
cdef load_buffer(bytes buf):
cdef load_ctx ctx = load_ctx()
ctx.offset = 1
ctx.buffer = buf
ctx.size = len(buf)
if len(buf) < 1:
raise NBTFormatError, "NBT Stream too short!"
if should_dump: print dump(buf)
assert ctx.buffer[0] == TAG_COMPOUND, "Data is not a TAG_Compound (found %d)" % ctx.buffer[0]
name = load_string(ctx)
#print "Root name", name
tag = load_compound(ctx)
tag.name = name
return tag
cdef load_byte(load_ctx ctx):
ctx.require(1)
cdef char * ptr = <char *>(ctx.buffer + ctx.offset)
ctx.offset += 1
cdef TAG_Byte tag = TAG_Byte.__new__(TAG_Byte)
tag.value = ptr[0]
tag.tagID = TAG_BYTE
tag._name = u""
return tag
cdef load_short(load_ctx ctx):
ctx.require(2)
cdef short * ptr = <short *>(ctx.buffer + ctx.offset)
swab(ptr, 2)
ctx.offset += 2
cdef TAG_Short tag = TAG_Short.__new__(TAG_Short)
tag.value = ptr[0]
tag.tagID = TAG_SHORT
tag._name = u""
return tag
cdef load_int(load_ctx ctx):
ctx.require(4)
cdef int * ptr = <int *>(ctx.buffer + ctx.offset)
swab(ptr, 4)
ctx.offset += 4
cdef TAG_Int tag = TAG_Int.__new__(TAG_Int)
tag.value = (ptr[0])
tag.tagID = TAG_INT
tag._name = u""
return tag
cdef load_long(load_ctx ctx):
ctx.require(8)
cdef long long * ptr = <long long *>(ctx.buffer + ctx.offset)
swab(ptr, 8)
ctx.offset += 8
cdef TAG_Long tag = TAG_Long.__new__(TAG_Long)
tag.value = ptr[0]
tag.tagID = TAG_LONG
tag._name = u""
return tag
cdef load_float(load_ctx ctx):
ctx.require(4)
cdef float * ptr = <float *>(ctx.buffer + ctx.offset)
swab(ptr, 4)
ctx.offset += 4
cdef TAG_Float tag = TAG_Float.__new__(TAG_Float)
tag.value = ptr[0]
tag.tagID = TAG_FLOAT
tag._name = u""
return tag
cdef load_double(load_ctx ctx):
ctx.require(8)
cdef double * ptr = <double *>(ctx.buffer + ctx.offset)
swab(ptr, 8)
ctx.offset += 8
cdef TAG_Double tag = TAG_Double.__new__(TAG_Double)
tag.value = ptr[0]
tag.tagID = TAG_DOUBLE
tag._name = u""
return tag
cdef load_bytearray(load_ctx ctx):
ctx.require(4)
cdef unsigned int * ptr = <unsigned int *>(ctx.buffer + ctx.offset)
swab(ptr, 4)
cdef unsigned int length = ptr[0]
ctx.offset += 4
cdef char * arr = ctx.buffer + ctx.offset
#print "Bytearray", length, ctx.size - ctx.offset
ctx.require(length)
ctx.offset += length
return TAG_Byte_Array(fromstring(arr[:length], dtype='uint8', count=length))
### --- load_compound ---
cdef load_compound(load_ctx ctx):
#print "load_compound buf=%d off=%d" % (ctx.buffer[0], ctx.offset)
cdef char tagID
cdef _TAG_Compound root_tag = TAG_Compound()
assert root_tag is not None
while True:
ctx.require(1)
tagID = ctx.buffer[ctx.offset]
ctx.offset += 1
if tagID == TAG_END:
#print "TAG_END at ", ctx.offset
break
else:
name = load_string(ctx)
tag = load_tag(tagID, ctx)
#tag.name = name
#print "tagID=%d name=%s at %d" % (tagID, tag.name, ctx.offset)
root_tag[name] = tag
return root_tag
cdef load_list(load_ctx ctx):
ctx.require(5)
cdef char tagID = ctx.buffer[ctx.offset]
ctx.offset += 1
cdef int * ptr = <int *>(ctx.buffer + ctx.offset)
swab(ptr, 4)
ctx.offset += 4
length = ptr[0]
cdef _TAG_List tag = TAG_List()
cdef list val = tag.value
cdef int i
for i in range(length):
PyList_Append(val, load_tag(tagID, ctx))
return tag
cdef unicode load_string(load_ctx ctx):
ctx.require(2)
cdef unsigned short * ptr = <unsigned short *>(ctx.buffer+ctx.offset)
swab(ptr, 2)
ctx.offset += 2
cdef unsigned short length = ptr[0]
#print "String: ", ctx.offset, length
cdef unicode u = PyUnicode_DecodeUTF8(ctx.buffer + ctx.offset, length, "strict")
ctx.offset += length
return u
cdef load_tag(char tagID, load_ctx ctx):
if tagID == TAG_BYTE:
return load_byte(ctx)
if tagID == TAG_SHORT:
return load_short(ctx)
if tagID == TAG_INT:
return load_int(ctx)
if tagID == TAG_LONG:
return load_long(ctx)
if tagID == TAG_FLOAT:
return load_float(ctx)
if tagID == TAG_DOUBLE:
return load_double(ctx)
if tagID == TAG_BYTE_ARRAY:
return load_bytearray(ctx)
if tagID == TAG_STRING:
u = load_string(ctx)
return TAG_String(u)
if tagID == TAG_LIST:
return load_list(ctx)
if tagID == TAG_COMPOUND:
return load_compound(ctx)
if tagID == TAG_INT_ARRAY:
pass
if tagID == TAG_SHORT_ARRAY:
pass
FILTER=''.join([(len(repr(chr(x)))==3) and chr(x) or '.' for x in range(256)])
def dump(src, length=8):
N=0; result=''
while src:
s,src = src[:length],src[length:]
hexa = ' '.join(["%02X"%ord(x) for x in s])
s = s.translate(FILTER)
result += "%04X %-*s %s\n" % (N, length*3, hexa, s)
N+=length
return result
cdef save_root_tag(tag, filename = "", buf = None):
sio = StringIO()
save_tag(tag, sio)
data = sio.getvalue()
if buf is None:
f = file(filename, "wb")
gzio = StringIO()
gz = gzip.GzipFile(fileobj=gzio, mode='wb', compresslevel=2)
gz.write(data)
gz.close()
f.write(gzio.getvalue())
else:
buf.write(data)
cdef save_tag(TAG_Value tag, object buf):
save_tag_id(tag.tagID, buf)
save_tag_name(tag, buf)
save_tag_value(tag, buf)
cdef save_tag_id(char tagID, object buf):
cwrite(buf, &tagID, 1)
cdef save_tag_name(TAG_Value tag, object buf):
name = tag.name.encode('utf-8')
save_string(name, buf)
cdef save_string(bytes value, object buf):
cdef unsigned short length = len(value)
cdef char * s = value
swab(&length, 2)
cwrite(buf, <char *>&length, 2)
cwrite(buf, s, len(value))
cdef save_byte_array(object value, object buf):
value = value.tostring()
cdef char * s = value
cdef unsigned int length = len(value)
swab(&length, 4)
cwrite(buf, <char *>&length, 4)
cwrite(buf, s, len(value))
cdef save_byte(char value, object buf):
cwrite(buf, <char *>&value, 1)
cdef save_short(short value, object buf):
swab(&value, 2)
cwrite(buf, <char *>&value, 2)
cdef save_int(int value, object buf):
swab(&value, 4)
cwrite(buf, <char *>&value, 4)
cdef save_long(long long value, object buf):
# print "Long long value: ", value, sizeof(value)
swab(&value, 8)
cdef char * p = <char *>&value
cwrite(buf, p, 8)
#cwrite(buf, p+4, 4)
#cwrite(buf, "\0\0\0\0\0\0\0\0", 8)
cdef save_float(float value, object buf):
swab(&value, 4)
cwrite(buf, <char *>&value, 4)
cdef save_double(double value, object buf):
swab(&value, 8)
cwrite(buf, <char *>&value, 8)
cdef save_tag_value(TAG_Value tag, object buf):
cdef char tagID = tag.tagID
if tagID == TAG_BYTE:
(<TAG_Byte>tag).save_value(buf)
if tagID == TAG_SHORT:
(<TAG_Short>tag).save_value(buf)
if tagID == TAG_INT:
(<TAG_Int>tag).save_value(buf)
if tagID == TAG_LONG:
(<TAG_Long>tag).save_value(buf)
if tagID == TAG_FLOAT:
(<TAG_Float>tag).save_value(buf)
if tagID == TAG_DOUBLE:
(<TAG_Double>tag).save_value(buf)
if tagID == TAG_BYTE_ARRAY:
(<TAG_Byte_Array>tag).save_value(buf)
if tagID == TAG_STRING:
(<TAG_String>tag).save_value(buf)
if tagID == TAG_LIST:
(<_TAG_List>tag).save_value(buf)
if tagID == TAG_COMPOUND:
(<_TAG_Compound>tag).save_value(buf)
if tagID == TAG_INT_ARRAY:
pass
if tagID == TAG_SHORT_ARRAY:
pass