scratch.py

import array
import itertools
import socket
import struct

class ScratchError(Exception): pass
class ScratchConnectionError(ScratchError): pass	

class Scratch(object):

    prefix_len = 4
    broadcast_prefix_len = prefix_len + len('broadcast ')
    sensorupdate_prefix_len = prefix_len + len('sensor-update ')

    msg_types = set(['broadcast', 'sensor-update'])

    def __init__(self, host='localhost', port=42001):
	self.host = host
	self.port = port
	self.socket = None
	self.connected = False
	self.connect()

    def __repr__(self):
	return "Scratch(host=%r, port=%r)" % (self.host, self.port)

    def _pack(self, msg):
	"""
	Packages msg according to Scratch message specification (encodes and 
	appends length prefix to msg). Credit to chalkmarrow from the 
	scratch.mit.edu forums for the prefix encoding code.
	"""
	n = len(msg) 
	a = array.array('c')
	a.append(chr((n >> 24) & 0xFF))
	a.append(chr((n >> 16) & 0xFF))
	a.append(chr((n >>  8) & 0xFF))
	a.append(chr(n & 0xFF))
	return a.tostring() + msg 

    def _extract_len(self, prefix):
	"""
	Extracts the length of a Scratch message from the given message prefix. 
	"""
	return struct.unpack(">L", prefix)[0]

    def _get_type(self, s):
	"""
	Converts a string from Scratch to its proper type in Python. Expects a
	string with its delimiting quotes in place. Returns either a string, 
	int or float. 
	"""
	# TODO: what if the number is bigger than an int or float?
	if s.startswith('"') and s.endswith('"'):
	    return s[1:-1]
	elif s.find('.') != -1: 
	    return float(s) 
	else: 
	    return int(s)

    def _escape(self, msg):
	"""
	Escapes double quotes by adding another double quote as per the Scratch
	protocol. Expects a string without its delimiting quotes. Returns a new
	escaped string. 
	"""
	escaped = ''	
	for c in msg: 
	    escaped += c
	    if c == '"':
		escaped += '"'
	return escaped

    def _unescape(self, msg):
	"""
	Removes double quotes that were used to escape double quotes. Expects
	a string without its delimiting quotes, or a number. Returns a new
	unescaped string.
	"""
	if isinstance(msg, (int, float, long)):
	    return msg

	unescaped = ''
	i = 0
	while i < len(msg):
	    unescaped += msg[i]
	    if msg[i] == '"':
		i+=1
	    i+=1
	return unescaped     

    def _is_msg(self, msg):
	"""
	Returns True if message is a proper Scratch message, else return False.
	"""
	if not msg or len(msg) < self.prefix_len:
	    return False
	length = self._extract_len(msg[:self.prefix_len])
	msg_type = msg[self.prefix_len:].split(' ', 1)[0]
	if length == len(msg[self.prefix_len:]) and msg_type in self.msg_types:
	    return True
	return False

    def _parse_broadcast(self, msg):
	"""
	Given a broacast message, returns the message that was broadcast.
	"""
	# get message, remove surrounding quotes, and unescape
	return self._unescape(self._get_type(msg[self.broadcast_prefix_len:]))

    def _parse_sensorupdate(self, msg):
	"""
	Given a sensor-update message, returns the sensors/variables that were
	updated as a dict that maps sensors/variables to their updated values.
	"""
	update = msg[self.sensorupdate_prefix_len:]
	parsed = [] # each element is either a sensor (key) or a sensor value
	curr_seg = '' # current segment (i.e. key or value) being formed
	numq = 0 # number of double quotes in current segment
	for seg in update.split(' ')[:-1]: # last char in update is a space
	    numq += seg.count('"')
	    curr_seg += seg
	    # even number of quotes means we've finished parsing a segment
	    if numq % 2 == 0: 
		parsed.append(curr_seg)
		curr_seg = ''
		numq = 0
	    else: # segment has a space inside, so add back it in
		curr_seg += ' '
	unescaped = [self._unescape(self._get_type(x)) for x in parsed]
	# combine into a dict using iterators (both elements in the list
	# inputted to izip have a reference to the same iterator). even 
	# elements are keys, odd are values
	return dict(itertools.izip(*[iter(unescaped)]*2)) 

    def _parse(self, msg):
	"""
	Parses a Scratch message and returns a tuple with the first element
	as the message type, and the second element as the message payload. The 
	payload for a 'broadcast' message is a string, and the payload for a 
	'sensor-update' message is a dict whose keys are variables, and values
	are updated variable values. Returns None if msg is not a message.
	"""
	if not self._is_msg(msg):
	    return None
	msg_type = msg[self.prefix_len:].split(' ')[0]
	if msg_type == 'broadcast':
	    return ('broadcast', self._parse_broadcast(msg))
	else:
	    return ('sensor-update', self._parse_sensorupdate(msg))

    def _write(self, data):
	"""
	Writes string data out to Scratch
	"""
	total_sent = 0
	length = len(data)
	while total_sent < length:
	    try:
		sent = self.socket.send(data[total_sent:])
	    except socket.error as (err, msg):
		self.connected = False
		raise ScratchError("[Errno %d] %s" % (err, msg))
	    if sent == 0:
		self.connected = False
		raise ScratchConnectionError("Connection broken")
	    total_sent += sent

    def _send(self, data):
	"""
	Sends a message to Scratch
	"""
	self._write(self._pack(data))

    def _read(self, size):
	"""
	Reads size number of bytes from Scratch and returns data as a string
	"""
	data = ''
	while len(data) < size:
	    try:
		chunk = self.socket.recv(size-len(data))
	    except socket.error as (err, msg):
		self.connected = False
		raise ScratchError("[Errno %d] %s" % (err, msg))
	    if chunk == '':
		self.connected = False
		raise ScratchConnectionError("Connection broken")
	    data += chunk
	return data

    def _recv(self):
	"""
	Receives and returns a message from Scratch
	"""
	prefix = self._read(self.prefix_len)
	msg = self._read(self._extract_len(prefix))
	return prefix + msg

    def connect(self):
	"""
	Connects to Scratch.
	"""
	self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
	try:
	    self.socket.connect((self.host, self.port))
	except socket.error as (err, msg):
	    self.connected = False
	    raise ScratchError("[Errno %d] %s" % (err, msg))
	self.connected = True

    def disconnect(self):
	"""
	Closes connection to Scratch
	"""
	try: # connection may already be disconnected, so catch exceptions
	    self.socket.shutdown(socket.SHUT_RDWR) # a proper disconnect
	except socket.error:
	    pass
	self.socket.close()
	self.connected = False

    def sensorupdate(self, data):
	"""
	Given a dict of sensors and values, updates those sensors with the 
	values in Scratch.
	"""
	if not isinstance(data, dict):
	    raise TypeError('Expected a dict')
	msg = 'sensor-update '
	for key in data.keys():
	    msg += '"%s" "%s" ' % (self._escape(str(key)), 
				   self._escape(str(data[key])))
	self._send(msg)

    def broadcast(self, msg):
	"""
	Broadcasts msg to Scratch. msg can be a single message or an iterable 
	(list, tuple, set, generator, etc.) of messages.
	"""
	if getattr(msg, '__iter__', False): # iterable
	    for m in msg:
		self._send('broadcast "%s"' % self._escape(str(m)))
	else: # probably a string or number
	    self._send('broadcast "%s"' % self._escape(str(msg)))

    def receive(self):
	"""
	Receives broadcasts and sensor updates from Scratch. Returns a tuple
	with the first element as the message type and the second element 
	as the message payload. broadcast messages have a string as payload, 
	and the sensor-update messages have a dict as payload. Returns None if 
	message received could not be parsed. Raises exceptions on connection
	errors.
	"""
	return self._parse(self._recv())