path: root/python/gats.py

#
# Copyright (C) 2007-2012 Xagasoft, All rights reserved.
#
# This file is part of the libgats library and is released under the
# terms of the license contained in the file LICENSE.
#

"""Provides pickle-compatible functions for working with Gats packets.
Converts to and from normal python structures like dictionaries, lists,
tuples, booleans, strings, etc.

This also provides a couple of helpers to make working with sockets a bit
easier, recv and send work with sockets in blocking mode."""

import struct
from cStringIO import StringIO as _StringIO
import math

_scalarMagic = float.fromhex('0x1.62e42fefa39efp+2')

def recv( s ):
	'''A helper that reads a single packet from the provided socket (s).  The
	socket should be in blocking mode for this to work correctly.  This function
	doesn't handle any exceptions, so if the socket closes while reading
	handling the exception is up to you.'''
	while True:
		version = s.recv(1)
		if len(version) == 0:
			return None
		version = ord(version[0])
		if version != 0:
			break

	if version == 1:
		sBuf = _StringIO()
		iSize = 5
		buf = ''
		while len(buf) < 4:
			buf += s.recv(4-len(buf))
		iGoalSize = struct.unpack('>I', buf )[0]
		while iSize < iGoalSize:
			buf = s.recv( min(iGoalSize-iSize, 4096) )
			sBuf.write( buf )
			iSize += len(buf)

		sBuf.seek( 0 )
		return _readObj( sBuf )

def send( obj, s ):
	'''A helper that sends the entire object (obj), in a Gats encoded packet to
	the provided socket (s).  This buffers the data first, so it should use
	as few network packets as possible to send the data.'''
	s.send( dumps( obj ) )

def loads( sIn ):
	'''Loads a complete Gats packet as described by the string sIn and returns
	the contained object.  The string can only contain one packet.  Any data
	after the first packet will simply be ignored.'''
	return load( _StringIO( sIn ) )

def dumps( obj ):
	'''Returns the provided object (obj) as a complete Gats packet in a string,
	since Gats encoded objects are not in plain text, it'll probably be useless
	to print these strings out.'''
	sTmp = _StringIO()
	dump( obj, sTmp )
	return sTmp.getvalue()

def load( sIn ):
	'''Reads a single Gats packet (one object) from the provided file-like
	object and returns it.  If it can't read anything it returns None.'''
	# Scan for a valid packet header
	while True:
		version = sIn.read(1)
		if len(version) == 0:
			return None
		version = ord(version[0])
		if version != 0:
			break

	if version == 1:
		size = struct.unpack('>I', sIn.read(4) )[0]
		return _readObj( sIn )

def dump( obj, sOut ):
	'''Writes the given object (obj) to the file-like object sOut as a complete
	Gats packet.'''
	sTmp = _StringIO()
	_writeObj( obj, sTmp )
	sCore = sTmp.getvalue()
	sOut.write(	struct.pack('>BI', 1, len(sCore)+5 ) )
	sOut.write( sCore )

class EndMarker:
	pass

def _readObj( sIn ):
	t = sIn.read( 1 )
	if t == 'i':	# Integer
		return _readPackedInt( sIn )
	elif t == 's':	# String
		lng = _readPackedInt( sIn )
		return sIn.read( lng )
	elif t == '0':	# Boolean false
		return False
	elif t == '1':	# Boolean true
		return True
	elif t == 'l':	# List
		ret = []
		while True:
			value = _readObj( sIn )
			if isinstance( value, EndMarker ):
				return ret
			ret.append( value )
	elif t == 'd':	# Dictionary
		ret = {}
		while True:
			key = _readObj( sIn )
			if isinstance( key, EndMarker ):
				return ret
			if not isinstance( key, str ):
				raise Exception('Only strings can be used as keys in gats dictionaries')
			value = _readObj( sIn );
			ret[key] = value
	elif t == 'f':	# Float
		pLng = _readPackedInt( sIn )
		bNeg = False
		if pLng < 0:
			bNeg = True
			pLng = -pLng

		dValue = 0.0
		dat = sIn.read( pLng )
		for i in xrange(len(dat)-1,0,-1):
			dValue = (dValue+ord(dat[i]))*(1.0/256.0)
		dValue += ord(dat[0])
		iScale = _readPackedInt( sIn )
		dValue *= pow( 256.0, iScale )
		if bNeg:
			return -dValue
		return dValue
	elif t == 'F':	# Exceptional float
		st = sIn.read(1)
		if st == 'N':
			return float('-nan')
		elif st == 'n':
			return float('nan')
		elif st == 'I':
			return float('-inf')
		elif st == 'i':
			return float('inf')
		elif st == 'Z':
			return -0.0
		elif st == 'z':
			return 0.0
		else:
			raise Exception('Invalid exceptional float subtype found.')
	elif t == 'n':
		return None
	elif t == 'e':	# End marker
		return EndMarker()
	else:
		raise Exception('Invalid gats type discovered: ' + t)
	return 'not implemented yet';

def _writeObj( obj, sOut ):
	if obj is None:
		sOut.write('n')
	elif isinstance( obj, bool ):
		if obj == True:
			sOut.write('1')
		else:
			sOut.write('0')
	elif isinstance( obj, int ):
		sOut.write('i')
		_writePackedInt( obj, sOut )
	elif isinstance( obj, str ):
		sOut.write('s')
		_writePackedInt( len(obj), sOut )
		sOut.write( obj )
	elif isinstance( obj, list ) or isinstance( obj, tuple ):
		sOut.write('l')
		for s in obj:
			_writeObj( s, sOut )
		sOut.write('e')
	elif isinstance( obj, dict ):
		sOut.write('d')
		for key in obj.iterkeys():
			_writeObj( key, sOut )
			_writeObj( obj[key], sOut )
		sOut.write('e')
	elif isinstance( obj, float ):
		if math.isnan( obj ):
			if math.copysign( 1.0, obj ) < 0.0:
				sOut.write('FN')
			else:
				sOut.write('Fn')
		elif math.isinf( obj ):
			if math.copysign( 1.0, obj ) < 0.0:
				sOut.write('FI')
			else:
				sOut.write('Fi')
		elif obj == 0.0:
			if math.copysign( 1.0, obj ) < 0.0:
				sOut.write('FZ')
			else:
				sOut.write('Fz')
		else:
			sOut.write('f')
			d = obj
			bNeg = False
			if d < 0.0:
				bNeg = True
				d = -d

			iScale = int(math.log( d ) / _scalarMagic)
			if iScale < 0:
				iScale -= 1

			sTmp = _StringIO()

			d /= pow( 256.0, iScale )
			sTmp.write( chr( int(d) ) )
			d -= int(d)
			for j in xrange( 0, 150 ):
				d = d*256.0
				sTmp.write( chr(int(d)) )
				d -= int(d)
				if d == 0.0:
					break

			sTmp = sTmp.getvalue()
			if bNeg:
				_writePackedInt( -len(sTmp), sOut )
			else:
				_writePackedInt( len(sTmp), sOut )
			sOut.write( sTmp )
			_writePackedInt( iScale, sOut )
	else:
		raise Exception('A type that is not gats-encodable was encountered: ' + str(type(obj)))

def _readPackedInt( sIn ):
	'''Internal helper function that reads an integer in packed format from the
	provided file-like object and returns it.'''
	bNeg = False

	b = ord(sIn.read(1)[0])
	bNeg = (b&0x40) == 0x40;
	rOut = b&0x3F;

	c = 0
	while (b&0x80) == 0x80:
		b = ord(sIn.read(1)[0])
		rOut |= (b&0x7F)<<(6+7*c)
		c += 1
	if bNeg:
		return -rOut
	return rOut

def _writePackedInt( iIn, sOut ):
	'''Internal helper function that writes an integer in packed format to the
	provided file-like object.'''
	if iIn < 0:
		iIn = -iIn
		b = iIn&0x3F
		if iIn > b:
			b |= 0x80 | 0x40
		else:
			b |= 0x40
	else:
		b = iIn&0x3F
		if iIn > b:
			b |= 0x80
	sOut.write( chr( b ) )

	iIn = iIn >> 6

	while iIn > 0:
		b = iIn&0x7F
		if iIn > b:
			b |= 0x80
		sOut.write( chr( b ) )
		iIn = iIn >> 7