def test_tuple_3_types(self):
types = Types()
typ = types.tuple_type(
types.float_type, types.uint64_type, types.string_type)
self.assertTrue(isinstance(typ, TupleType))
self.assertEqual(typ.python_type, tuple)
self.check_protobuf_type(
Type.TUPLE, '', '', 3, typ.protobuf_type)
self.check_protobuf_type(
Type.FLOAT, '', '', 0, typ.protobuf_type.types[0])
self.check_protobuf_type(
Type.UINT64, '', '', 0, typ.protobuf_type.types[1])
self.check_protobuf_type(
Type.STRING, '', '', 0, typ.protobuf_type.types[2])
self.assertEqual(3, len(typ.value_types))
self.assertTrue(isinstance(typ.value_types[0], ValueType))
self.assertTrue(isinstance(typ.value_types[1], ValueType))
self.assertTrue(isinstance(typ.value_types[2], ValueType))
self.assertEqual(float, typ.value_types[0].python_type)
self.assertEqual(long, typ.value_types[1].python_type)
self.assertEqual(str, typ.value_types[2].python_type)
self.check_protobuf_type(
Type.FLOAT, '', '', 0, typ.value_types[0].protobuf_type)
self.check_protobuf_type(
Type.UINT64, '', '', 0, typ.value_types[1].protobuf_type)
self.check_protobuf_type(
Type.STRING, '', '', 0, typ.value_types[2].protobuf_type)
def test_get_return_type(self):
types = Types()
self.assertEqual('float', types.get_return_type('float', []).protobuf_type)
self.assertEqual('int32', types.get_return_type('int32', []).protobuf_type)
self.assertEqual('KRPC.Response', types.get_return_type('KRPC.Response', []).protobuf_type)
self.assertEqual('Class(ServiceName.ClassName)',
types.get_return_type('uint64', ['ReturnType.Class(ServiceName.ClassName)']).protobuf_type)
def test_coerce_to(self):
types = Types()
cases = [
(42.0, 42, 'double'),
(42.0, 42, 'float'),
(42, 42.0, 'int32'),
(42, 42L, 'int32'),
(42L, 42.0, 'int64'),
(42L, 42, 'int64'),
(42, 42.0, 'uint32'),
(42, 42L, 'uint32'),
(42L, 42.0, 'uint64'),
(42L, 42, 'uint64'),
(list(), tuple(), 'List(string)'),
((0,1,2), [0,1,2], 'Tuple(int32,int32,int32)'),
([0,1,2], (0,1,2), 'List(int32)'),
]
for expected, value, typ in cases:
coerced_value = types.coerce_to(value, types.as_type(typ))
self.assertEqual(expected, coerced_value)
self.assertEqual(type(expected), type(coerced_value))
self.assertRaises(ValueError, types.coerce_to, None, types.as_type('float'))
self.assertRaises(ValueError, types.coerce_to, '', types.as_type('float'))
self.assertRaises(ValueError, types.coerce_to, True, types.as_type('float'))
self.assertRaises(ValueError, types.coerce_to, list(), types.as_type('Tuple(int32)'))
self.assertRaises(ValueError, types.coerce_to, ["foo",2], types.as_type('Tuple(string)'))
self.assertRaises(ValueError, types.coerce_to, [1], types.as_type('Tuple(string)'))
self.assertRaises(ValueError, types.coerce_to, [1,"a","b"], types.as_type('List(string)'))
def test_value_types(self):
types = Types()
for protobuf_typ in PROTOBUF_VALUE_TYPES:
python_typ = PROTOBUF_TO_PYTHON_VALUE_TYPE[protobuf_typ]
typ = types.as_type(protobuf_typ)
self.assertTrue(isinstance(typ, ValueType))
self.assertEqual(protobuf_typ, typ.protobuf_type)
self.assertEqual(python_typ, typ.python_type)
self.assertRaises(ValueError, ValueType, 'invalid')
def test_protobuf_enum_types(self):
types = Types()
typ = types.as_type('Test.TestEnum')
self.assertTrue(isinstance(typ, ProtobufEnumType))
self.assertEqual(int, typ.python_type)
self.assertEqual('Test.TestEnum', typ.protobuf_type)
self.assertRaises(ValueError, types.as_type, 'Test.DoesntExist')
self.assertRaises(ValueError, ProtobufEnumType, '')
self.assertRaises(ValueError, ProtobufEnumType, 'invalid')
self.assertRaises(ValueError, ProtobufEnumType, '.')
self.assertRaises(ValueError, ProtobufEnumType, 'foo.bar')
def test_message_types(self):
types = Types()
typ = types.as_type('KRPC.Request')
self.assertTrue(isinstance(typ, MessageType))
self.assertEqual(krpc.schema.KRPC.Request, typ.python_type)
self.assertEqual('KRPC.Request', typ.protobuf_type)
self.assertRaises(ValueError, types.as_type, 'KRPC.DoesntExist')
self.assertRaises(ValueError, MessageType, '')
self.assertRaises(ValueError, MessageType, 'invalid')
self.assertRaises(ValueError, MessageType, '.')
self.assertRaises(ValueError, MessageType, 'foo.bar')
def test_enum_types(self):
types = Types()
typ = types.as_type('Enum(ServiceName.EnumName)')
self.assertTrue(isinstance(typ, EnumType))
self.assertEqual(None, typ.python_type)
self.assertTrue('Enum(ServiceName.EnumName)', typ.protobuf_type)
typ.set_values({'a': 0, 'b': 42, 'c': 100})
self.assertTrue(issubclass(typ.python_type, Enum))
self.assertEquals(0, typ.python_type.a.value)
self.assertEquals(42, typ.python_type.b.value)
self.assertEquals(100, typ.python_type.c.value)
def test_class_types(self):
types = Types()
typ = types.as_type('Class(ServiceName.ClassName)')
self.assertTrue(isinstance(typ, ClassType))
self.assertTrue(issubclass(typ.python_type, ClassBase))
self.assertEqual('Class(ServiceName.ClassName)', typ.protobuf_type)
instance = typ.python_type(42)
self.assertEqual(42, instance._object_id)
self.assertEqual('ServiceName', instance._service_name)
self.assertEqual('ClassName', instance._class_name)
typ2 = types.as_type('Class(ServiceName.ClassName)')
self.assertEqual(typ, typ2)
def test_set_types(self):
types = Types()
typ = types.set_type(types.string_type)
self.assertTrue(isinstance(typ, SetType))
self.assertEqual(typ.python_type, set)
self.check_protobuf_type(
Type.SET, '', '', 1, typ.protobuf_type)
self.check_protobuf_type(
Type.STRING, '', '', 0, typ.protobuf_type.types[0])
self.assertTrue(isinstance(typ.value_type, ValueType))
self.assertEqual(str, typ.value_type.python_type)
self.check_protobuf_type(
Type.STRING, '', '', 0, typ.value_type.protobuf_type)
def test_list_types(self):
types = Types()
typ = types.list_type(types.uint32_type)
self.assertTrue(isinstance(typ, ListType))
self.assertEqual(typ.python_type, list)
self.check_protobuf_type(
Type.LIST, '', '', 1, typ.protobuf_type)
self.check_protobuf_type(
Type.UINT32, '', '', 0, typ.protobuf_type.types[0])
self.assertTrue(isinstance(typ.value_type, ValueType))
self.assertEqual(int, typ.value_type.python_type)
self.check_protobuf_type(
Type.UINT32, '', '', 0, typ.value_type.protobuf_type)
def test_coerce_to(self):
types = Types()
cases = [
(42.0, 42, types.double_type),
(42.0, 42, types.float_type),
(42, 42.0, types.sint32_type),
(42, 42L, types.sint32_type),
(42L, 42.0, types.sint64_type),
(42L, 42, types.sint64_type),
(42, 42.0, types.uint32_type),
(42, 42L, types.uint32_type),
(42L, 42.0, types.uint64_type),
(42L, 42, types.uint64_type),
(list(), tuple(), types.list_type(types.string_type)),
((0, 1, 2), [0, 1, 2],
types.tuple_type(types.sint32_type,
types.sint32_type,
types.sint32_type)),
([0, 1, 2], (0, 1, 2),
types.list_type(types.sint32_type)),
(['foo', 'bar'], ['foo', 'bar'],
types.list_type(types.string_type))
]
for expected, value, typ in cases:
coerced_value = types.coerce_to(value, typ)
self.assertEqual(expected, coerced_value)
self.assertEqual(type(expected), type(coerced_value))
strings = [
u'foo',
u'\xe2\x84\xa2',
u'Mystery Goo\xe2\x84\xa2 Containment Unit'
]
for string in strings:
self.assertEqual(
string, types.coerce_to(string, types.string_type))
self.assertRaises(ValueError, types.coerce_to,
None, types.float_type)
self.assertRaises(ValueError, types.coerce_to,
'', types.float_type)
self.assertRaises(ValueError, types.coerce_to,
True, types.float_type)
self.assertRaises(ValueError, types.coerce_to,
list(), types.tuple_type(types.uint32_type))
self.assertRaises(ValueError, types.coerce_to,
['foo', 2], types.tuple_type(types.string_type))
self.assertRaises(ValueError, types.coerce_to,
[1], types.tuple_type(types.string_type))
def test_tuple_1_types(self):
types = Types()
typ = types.tuple_type(types.bool_type)
self.assertTrue(isinstance(typ, TupleType))
self.assertEqual(typ.python_type, tuple)
self.check_protobuf_type(
Type.TUPLE, '', '', 1, typ.protobuf_type)
self.check_protobuf_type(
Type.BOOL, '', '', 0, typ.protobuf_type.types[0])
self.assertEqual(1, len(typ.value_types))
self.assertTrue(isinstance(typ.value_types[0], ValueType))
self.assertEqual(bool, typ.value_types[0].python_type)
self.check_protobuf_type(
Type.BOOL, '', '', 0, typ.value_types[0].protobuf_type)
def test_list_types(self):
types = Types()
typ = types.as_type('List(int32)')
self.assertTrue(isinstance(typ, ListType))
self.assertEqual(typ.python_type, list)
self.assertEqual('List(int32)', typ.protobuf_type)
self.assertTrue(isinstance(typ.value_type, ValueType))
self.assertEqual('int32', typ.value_type.protobuf_type)
self.assertEqual(int, typ.value_type.python_type)
self.assertRaises(ValueError, types.as_type, 'List')
self.assertRaises(ValueError, types.as_type, 'List(')
self.assertRaises(ValueError, types.as_type, 'List()')
self.assertRaises(ValueError, types.as_type, 'List(foo')
self.assertRaises(ValueError, types.as_type, 'List(int32,string)')
def test_class_types(self):
types = Types()
typ = types.class_type(
'ServiceName', 'ClassName', 'class documentation')
self.assertTrue(isinstance(typ, ClassType))
self.assertTrue(issubclass(typ.python_type, ClassBase))
self.assertEqual('class documentation', typ.python_type.__doc__)
self.check_protobuf_type(
Type.CLASS, 'ServiceName', 'ClassName', 0, typ.protobuf_type)
instance = typ.python_type(42)
self.assertEqual(42, instance._object_id)
self.assertEqual('ServiceName', instance._service_name)
self.assertEqual('ClassName', instance._class_name)
typ2 = types.as_type(typ.protobuf_type)
self.assertEqual(typ, typ2)
def __init__(self, rpc_connection, stream_connection):
self._types = Types()
self._rpc_connection = rpc_connection
self._rpc_connection_lock = threading.Lock()
self._stream_connection = stream_connection
self._stream_cache = {}
self._stream_cache_lock = threading.Lock()
self._request_type = self._types.as_type('KRPC.Request')
self._response_type = self._types.as_type('KRPC.Response')
# Get the services
services = self._invoke('KRPC', 'GetServices', [], [], [], self._types.as_type('KRPC.Services')).services
# Set up services
for service in services:
setattr(self, snake_case(service.name), create_service(self, service))
# Set up stream update thread
if stream_connection is not None:
self._stream_thread_stop = threading.Event()
self._stream_thread = threading.Thread(target=krpc.stream.update_thread,
args=(stream_connection, self._stream_thread_stop,
self._stream_cache, self._stream_cache_lock))
self._stream_thread.daemon = True
self._stream_thread.start()
else:
self._stream_thread = None
def __init__(self, rpc_connection, stream_connection):
self._types = Types()
self._rpc_connection = rpc_connection
self._rpc_connection_lock = threading.Lock()
self._stream_connection = stream_connection
self._request_type = self._types.as_type("KRPC.Request")
self._response_type = self._types.as_type("KRPC.Response")
# Set up the main KRPC service
self.krpc = KRPC(self)
services = self.krpc.get_services().services
# Create class types
# TODO: is this needed?!?
for service in services:
for procedure in service.procedures:
try:
name = Attributes.get_class_name(procedure.attributes)
self._types.as_type("Class(" + service.name + "." + name + ")")
except ValueError:
pass
# Set up services
for service in services:
if service.name != "KRPC":
setattr(self, _to_snake_case(service.name), create_service(self, service))
# Set up stream update thread
if stream_connection is not None:
self._stream_thread = threading.Thread(target=krpc.stream.update_thread, args=(stream_connection,))
self._stream_thread.daemon = True
self._stream_thread.start()
else:
self._stream_thread = None
def test_set_types(self):
types = Types()
typ = types.as_type('Set(string)')
self.assertTrue(isinstance(typ, SetType))
self.assertEqual(typ.python_type, set)
self.assertEqual('Set(string)', typ.protobuf_type)
self.assertTrue(isinstance(typ.value_type, ValueType))
self.assertEqual('string', typ.value_type.protobuf_type)
self.assertEqual(str, typ.value_type.python_type)
self.assertRaises(ValueError, types.as_type, 'Set')
self.assertRaises(ValueError, types.as_type, 'Set(')
self.assertRaises(ValueError, types.as_type, 'Set()')
self.assertRaises(ValueError, types.as_type, 'Set(string,)')
self.assertRaises(ValueError, types.as_type, 'Set(,)')
self.assertRaises(ValueError, types.as_type, 'Set(,string)')
self.assertRaises(ValueError, types.as_type, 'Set(int,string))')
def test_dictionary_types(self):
types = Types()
typ = types.dictionary_type(types.string_type, types.uint32_type)
self.assertTrue(isinstance(typ, DictionaryType))
self.assertEqual(typ.python_type, dict)
self.check_protobuf_type(
Type.DICTIONARY, '', '', 2, typ.protobuf_type)
self.check_protobuf_type(
Type.STRING, '', '', 0, typ.protobuf_type.types[0])
self.check_protobuf_type(
Type.UINT32, '', '', 0, typ.protobuf_type.types[1])
self.assertTrue(isinstance(typ.key_type, ValueType))
self.assertEqual(str, typ.key_type.python_type)
self.check_protobuf_type(
Type.STRING, '', '', 0, typ.key_type.protobuf_type)
self.assertTrue(isinstance(typ.value_type, ValueType))
self.assertEqual(int, typ.value_type.python_type)
self.check_protobuf_type(
Type.UINT32, '', '', 0, typ.value_type.protobuf_type)
def test_get_parameter_type(self):
types = Types()
self.assertEqual(float, types.get_parameter_type(0, 'float', []).python_type)
self.assertEqual('int32', types.get_parameter_type(0, 'int32', []).protobuf_type)
self.assertEqual('KRPC.Response', types.get_parameter_type(1, 'KRPC.Response', []).protobuf_type)
class_parameter = types.get_parameter_type(0, 'uint64', ['ParameterType(0).Class(ServiceName.ClassName)'])
self.assertEqual(types.as_type('Class(ServiceName.ClassName)'), class_parameter)
self.assertTrue(isinstance(class_parameter, ClassType))
self.assertTrue(issubclass(class_parameter.python_type, ClassBase))
self.assertEqual('Class(ServiceName.ClassName)', class_parameter.protobuf_type)
self.assertEqual('uint64', types.get_parameter_type(0, 'uint64', ['ParameterType(1).Class(ServiceName.ClassName)']).protobuf_type)
self.assertEqual('Test.TestEnum', types.get_parameter_type(0, 'Test.TestEnum', []).protobuf_type)
def test_tuple_types(self):
types = Types()
typ = types.as_type('Tuple(bool)')
self.assertTrue(isinstance(typ, TupleType))
self.assertEqual(typ.python_type, tuple)
self.assertEqual('Tuple(bool)', typ.protobuf_type)
self.assertEqual(1, len(typ.value_types))
self.assertTrue(isinstance(typ.value_types[0], ValueType))
self.assertEqual('bool', typ.value_types[0].protobuf_type)
self.assertEqual(bool, typ.value_types[0].python_type)
typ = types.as_type('Tuple(int32,string)')
self.assertTrue(isinstance(typ, TupleType))
self.assertEqual(typ.python_type, tuple)
self.assertEqual('Tuple(int32,string)', typ.protobuf_type)
self.assertEqual(2, len(typ.value_types))
self.assertTrue(isinstance(typ.value_types[0], ValueType))
self.assertTrue(isinstance(typ.value_types[1], ValueType))
self.assertEqual('int32', typ.value_types[0].protobuf_type)
self.assertEqual('string', typ.value_types[1].protobuf_type)
self.assertEqual(int, typ.value_types[0].python_type)
self.assertEqual(str, typ.value_types[1].python_type)
typ = types.as_type('Tuple(float,int64,string)')
self.assertTrue(isinstance(typ, TupleType))
self.assertEqual(typ.python_type, tuple)
self.assertEqual('Tuple(float,int64,string)', typ.protobuf_type)
self.assertEqual(3, len(typ.value_types))
self.assertTrue(isinstance(typ.value_types[0], ValueType))
self.assertTrue(isinstance(typ.value_types[1], ValueType))
self.assertTrue(isinstance(typ.value_types[2], ValueType))
self.assertEqual('float', typ.value_types[0].protobuf_type)
self.assertEqual('int64', typ.value_types[1].protobuf_type)
self.assertEqual('string', typ.value_types[2].protobuf_type)
self.assertEqual(float, typ.value_types[0].python_type)
self.assertEqual(long, typ.value_types[1].python_type)
self.assertEqual(str, typ.value_types[2].python_type)
self.assertRaises(ValueError, types.as_type, 'Tuple')
self.assertRaises(ValueError, types.as_type, 'Tuple(')
self.assertRaises(ValueError, types.as_type, 'Tuple()')
self.assertRaises(ValueError, types.as_type, 'Tuple(foo')
self.assertRaises(ValueError, types.as_type, 'Tuple(string,)')
self.assertRaises(ValueError, types.as_type, 'Tuple(,)')
self.assertRaises(ValueError, types.as_type, 'Tuple(,string)')
self.assertRaises(ValueError, types.as_type, 'Tuple(int,string))')
def test_tuple_2_types(self):
types = Types()
typ = types.tuple_type(types.uint32_type, types.string_type)
self.assertTrue(isinstance(typ, TupleType))
self.assertEqual(typ.python_type, tuple)
self.check_protobuf_type(
Type.TUPLE, '', '', 2, typ.protobuf_type)
self.check_protobuf_type(
Type.UINT32, '', '', 0, typ.protobuf_type.types[0])
self.check_protobuf_type(
Type.STRING, '', '', 0, typ.protobuf_type.types[1])
self.assertEqual(2, len(typ.value_types))
self.assertTrue(isinstance(typ.value_types[0], ValueType))
self.assertTrue(isinstance(typ.value_types[1], ValueType))
self.assertEqual(int, typ.value_types[0].python_type)
self.assertEqual(str, typ.value_types[1].python_type)
self.check_protobuf_type(
Type.UINT32, '', '', 0, typ.value_types[0].protobuf_type)
self.check_protobuf_type(
Type.STRING, '', '', 0, typ.value_types[1].protobuf_type)
def test_dictionary_types(self):
types = Types()
typ = types.as_type('Dictionary(string,int32)')
self.assertTrue(isinstance(typ, DictionaryType))
self.assertEqual(typ.python_type, dict)
self.assertEqual('Dictionary(string,int32)', typ.protobuf_type)
self.assertTrue(isinstance(typ.key_type, ValueType))
self.assertEqual('string', typ.key_type.protobuf_type)
self.assertEqual(str, typ.key_type.python_type)
self.assertTrue(isinstance(typ.value_type, ValueType))
self.assertEqual('int32', typ.value_type.protobuf_type)
self.assertEqual(int, typ.value_type.python_type)
self.assertRaises(ValueError, types.as_type, 'Dictionary')
self.assertRaises(ValueError, types.as_type, 'Dictionary(')
self.assertRaises(ValueError, types.as_type, 'Dictionary()')
self.assertRaises(ValueError, types.as_type, 'Dictionary(foo')
self.assertRaises(ValueError, types.as_type, 'Dictionary(string)')
self.assertRaises(ValueError, types.as_type, 'Dictionary(string,)')
self.assertRaises(ValueError, types.as_type, 'Dictionary(,)')
self.assertRaises(ValueError, types.as_type, 'Dictionary(,string)')
self.assertRaises(ValueError, types.as_type, 'Dictionary(int,string))')
def test_enumeration_types(self):
types = Types()
typ = types.enumeration_type(
'ServiceName', 'EnumName', 'enum documentation')
self.assertTrue(isinstance(typ, EnumerationType))
self.assertIsNone(typ.python_type)
self.check_protobuf_type(
Type.ENUMERATION, 'ServiceName', 'EnumName', 0, typ.protobuf_type)
typ.set_values({
'a': {'value': 0, 'doc': 'doca'},
'b': {'value': 42, 'doc': 'docb'},
'c': {'value': 100, 'doc': 'docc'}
})
self.assertTrue(issubclass(typ.python_type, Enum))
self.assertEqual('enum documentation', typ.python_type.__doc__)
self.assertEquals(0, typ.python_type.a.value)
self.assertEquals(42, typ.python_type.b.value)
self.assertEquals(100, typ.python_type.c.value)
self.assertEquals('doca', typ.python_type.a.__doc__)
self.assertEquals('docb', typ.python_type.b.__doc__)
self.assertEquals('docc', typ.python_type.c.__doc__)
typ2 = types.as_type(typ.protobuf_type)
self.assertEqual(typ, typ2)
def _parse_procedure(cls, procedure):
param_names = [snake_case(param.name)
for param in procedure.parameters]
param_types = [cls._client._types.as_type(param.type)
for param in procedure.parameters]
param_required = [not param.default_value
for param in procedure.parameters]
param_default = []
for param, typ in zip(procedure.parameters, param_types):
if param.default_value:
param_default.append(Decoder.decode(param.default_value, typ))
else:
param_default.append(None)
return_type = None
if not Types.is_none_type(procedure.return_type):
return_type = cls._client._types.as_type(procedure.return_type)
return param_names, param_types, param_required, \
param_default, return_type
def __init__(self, rpc_connection, stream_connection):
self._types = Types()
self._rpc_connection = rpc_connection
self._rpc_connection_lock = threading.Lock()
self._stream_connection = stream_connection
self._request_type = self._types.as_type('KRPC.Request')
self._response_type = self._types.as_type('KRPC.Response')
# Set up the main KRPC service
self.krpc = KRPC(self)
services = self.krpc.get_services().services
# Create class types
#TODO: is this needed?!?
#for service in services:
# for procedure in service.procedures:
# try:
# name = Attributes.get_class_name(procedure.attributes)
# self._types.as_type('Class(' + service.name + '.' + name + ')')
# except ValueError:
# pass
# Set up services
for service in services:
if service.name != 'KRPC':
setattr(self, snake_case(service.name), create_service(self, service))
# Set up stream update thread
if stream_connection is not None:
self._stream_thread_stop = threading.Event()
self._stream_thread = threading.Thread(target=krpc.stream.update_thread,
args=(stream_connection,self._stream_thread_stop))
self._stream_thread.daemon = True
self._stream_thread.start()
else:
self._stream_thread = None
class Client(object):
"""
A kRPC client, through which all Remote Procedure Calls are made.
Services provided by the server that the client connects to are automatically added.
RPCs can be made using client.ServiceName.ProcedureName(parameter)
"""
def __init__(self, rpc_connection, stream_connection):
self._types = Types()
self._rpc_connection = rpc_connection
self._rpc_connection_lock = threading.Lock()
self._stream_connection = stream_connection
self._request_type = self._types.as_type("KRPC.Request")
self._response_type = self._types.as_type("KRPC.Response")
# Set up the main KRPC service
self.krpc = KRPC(self)
services = self.krpc.get_services().services
# Create class types
# TODO: is this needed?!?
for service in services:
for procedure in service.procedures:
try:
name = Attributes.get_class_name(procedure.attributes)
self._types.as_type("Class(" + service.name + "." + name + ")")
except ValueError:
pass
# Set up services
for service in services:
if service.name != "KRPC":
setattr(self, _to_snake_case(service.name), create_service(self, service))
# Set up stream update thread
if stream_connection is not None:
self._stream_thread = threading.Thread(target=krpc.stream.update_thread, args=(stream_connection,))
self._stream_thread.daemon = True
self._stream_thread.start()
else:
self._stream_thread = None
def close(self):
self._rpc_connection.close()
if self._stream_connection is not None:
self._stream_connection.close()
def __enter__(self):
return self
def __exit__(self, typ, value, traceback):
self.close()
def add_stream(self, func, *args, **kwargs):
if self._stream_connection is None:
raise RuntimeError("Not connected to stream server")
return krpc.stream.add_stream(self, func, *args, **kwargs)
@contextmanager
def stream(self, func, *args, **kwargs):
""" 'with' support """
s = self.add_stream(func, *args, **kwargs)
try:
yield s
finally:
s.remove()
def _invoke(self, service, procedure, args=[], kwargs={}, param_names=[], param_types=[], return_type=None):
""" Execute an RPC """
# Build the request
request = self._build_request(service, procedure, args, kwargs, param_names, param_types, return_type)
# Send the request
with self._rpc_connection_lock:
self._send_request(request)
response = self._receive_response()
# Check for an error response
if response.HasField("error"):
raise RPCError(response.error)
# Decode the response and return the (optional) result
result = None
if return_type is not None:
result = Decoder.decode(response.return_value, return_type)
return result
def _build_request(self, service, procedure, args=[], kwargs={}, param_names=[], param_types=[], return_type=None):
""" Build a KRPC.Request object """
def encode_argument(i, value):
typ = param_types[i]
if type(value) != typ.python_type:
# Try coercing to the correct type
try:
value = self._types.coerce_to(value, typ)
except ValueError:
raise TypeError(
"%s.%s() argument %d must be a %s, got a %s"
% (service, procedure, i, typ.python_type, type(value))
)
return Encoder.encode(value, typ)
if len(args) > len(param_types):
raise TypeError(
"%s.%s() takes exactly %d arguments (%d given)" % (service, procedure, len(param_types), len(args))
)
arguments = []
nargs = len(args)
for i, param in enumerate(param_names):
add = False
if i < nargs and not isinstance(args[i], DefaultArgument):
arg = args[i]
add = True
elif param in kwargs:
arg = args[param]
add = True
if add:
argument = krpc.schema.KRPC.Argument()
argument.position = i
argument.value = encode_argument(i, arg)
arguments.append(argument)
# Build the request object
request = krpc.schema.KRPC.Request()
request.service = service
request.procedure = procedure
request.arguments.extend(arguments)
return request
def _send_request(self, request):
""" Send a KRPC.Request object to the server """
data = Encoder.encode_delimited(request, self._request_type)
self._rpc_connection.send(data)
def _receive_response(self):
""" Receive data from the server and decode it into a KRPC.Response object """
# Read the size and position of the response message
data = b""
while True:
try:
data += self._rpc_connection.partial_receive(1)
size, position = Decoder.decode_size_and_position(data)
break
except IndexError:
pass
# Read and decode the response message
data = self._rpc_connection.receive(size)
return Decoder.decode(data, self._response_type)
from collections import OrderedDict, defaultdict
from .docparser import DocumentationParser
from krpc.attributes import Attributes
from krpc.types import Types
from krpc.decoder import Decoder
types = Types()
class Appendable(object):
def __init__(self):
self._appended = []
def append(self, value):
self._appended.append(value)
@property
def appended(self):
return '\n\n'.join(self._appended)
class Service(Appendable):
def __init__(self, name, procedures, classes, enumerations, documentation,
sort):
super(Service, self).__init__()
self.name = name
self.fullname = name
self.documentation = documentation
self.cref = 'T:%s' % name
members = []
class Client(object):
"""
A kRPC client, through which all Remote Procedure Calls are made.
Services provided by the server that the client connects
to are automatically added. RPCs can be made using
client.ServiceName.ProcedureName(parameter)
"""
def __init__(self, rpc_connection, stream_connection):
self._types = Types()
self._rpc_connection = rpc_connection
self._rpc_connection_lock = threading.Lock()
self._stream_connection = stream_connection
self._stream_manager = StreamManager(self)
# Get the services
services = self._invoke('KRPC', 'GetServices', [], [], [],
self._types.services_type).services
# Set up services
for service in services:
setattr(self, snake_case(service.name),
create_service(self, service))
# Set up stream update thread
if stream_connection is not None:
self._stream_thread_stop = threading.Event()
self._stream_thread = threading.Thread(
target=krpc.streammanager.update_thread,
args=(self._stream_manager, stream_connection,
self._stream_thread_stop))
self._stream_thread.daemon = True
self._stream_thread.start()
else:
self._stream_thread = None
def close(self):
self._rpc_connection.close()
if self._stream_thread is not None:
self._stream_thread_stop.set()
self._stream_thread.join()
def __enter__(self):
return self
def __exit__(self, typ, value, traceback):
self.close()
def add_stream(self, func, *args, **kwargs):
""" Add a stream to the server """
if self._stream_connection is None:
raise StreamError('Not connected to stream server')
if func == setattr:
raise StreamError('Cannot stream a property setter')
return_type = self._get_return_type(func, *args, **kwargs)
call = self.get_call(func, *args, **kwargs)
return krpc.stream.Stream.from_call(self, return_type, call)
@contextmanager
def stream(self, func, *args, **kwargs):
""" 'with' support for add_stream """
stream = self.add_stream(func, *args, **kwargs)
try:
yield stream
finally:
stream.remove()
@property
def stream_update_condition(self):
""" Condition variable that is notified when
a stream update message has finished being processed. """
return self._stream_manager.update_condition
def wait_for_stream_update(self, timeout=None):
""" Wait until the next stream update message or a timeout occurs.
The condition variable must be locked before calling this method.
When timeout is not None, it should be a floating point number
specifying the timeout in seconds for the operation. """
self._stream_manager.wait_for_update(timeout)
def add_stream_update_callback(self, callback):
""" Add a callback that is invoked whenever
a stream update message has finished being processed. """
self._stream_manager.add_update_callback(callback)
def remove_stream_update_callback(self, callback):
""" Remove a stream update callback. """
self._stream_manager.remove_update_callback(callback)
@staticmethod
def get_call(func, *args, **kwargs):
""" Convert a remote procedure call to a KRPC.ProcedureCall message """
if func == getattr:
# A property or class property getter
attr = func(args[0].__class__, args[1])
return attr.fget._build_call(args[0])
elif func == setattr:
# A property setter
raise StreamError('Cannot create a call for a property setter')
elif hasattr(func, '__self__'):
# A method
return func._build_call(func.__self__, *args, **kwargs)
else:
# A class method
return func._build_call(*args, **kwargs)
@staticmethod
def _get_return_type(func, *args, **kwargs): # pylint: disable=unused-argument
""" Get the return type for a remote procedure call """
if func == getattr:
# A property or class property getter
attr = func(args[0].__class__, args[1])
return attr.fget._return_type
elif func == setattr:
# A property setter
raise StreamError('Cannot get return type for a property setter')
elif hasattr(func, '__self__'):
# A method
return func._return_type
else:
# A class method
return func._return_type
def _invoke(self, service, procedure, args, param_names, param_types,
return_type):
""" Execute an RPC """
# Build the request
call = self._build_call(service, procedure, args, param_names,
param_types, return_type)
request = KRPC.Request()
request.calls.extend([call])
# Send the request
with self._rpc_connection_lock:
self._rpc_connection.send_message(request)
response = self._rpc_connection.receive_message(KRPC.Response)
# Check for an error response
if response.HasField('error'):
raise self._build_error(response.error)
# Check for an error in the procedure results
if response.results[0].HasField('error'):
raise self._build_error(response.results[0].error)
# Decode the response and return the (optional) result
result = None
if return_type is not None:
result = Decoder.decode(response.results[0].value, return_type)
if isinstance(result, KRPC.Event):
result = Event(self, result)
return result
def _build_call(self, service, procedure, args, param_names, param_types,
return_type):
# pylint: disable=unused-argument
""" Build a KRPC.ProcedureCall object """
call = KRPC.ProcedureCall()
call.service = service
call.procedure = procedure
for i, (value, typ) in enumerate(zip(args, param_types)):
if isinstance(value, DefaultArgument):
continue
if not isinstance(value, typ.python_type):
try:
value = self._types.coerce_to(value, typ)
except ValueError:
raise TypeError(
'%s.%s() argument %d must be a %s, got a %s' %
(service, procedure, i, typ.python_type, type(value)))
call.arguments.add(position=i, value=Encoder.encode(value, typ))
return call
def _build_error(self, error):
""" Build an exception from an error message that
can be thrown to the calling code """
# TODO: modify the stack trace of the thrown exception so it looks like
# it came from the local call
if error.service and error.name:
service_name = snake_case(error.service)
type_name = error.name
if not hasattr(self, service_name):
raise RuntimeError(
'Error building exception; service \'%s\' not found' %
service_name)
service = getattr(self, service_name)
if not hasattr(service, type_name):
raise RuntimeError(
'Error building exception; type \'%s.%s\' not found' %
(service_name, type_name))
return getattr(service, type_name)(self._error_message(error))
return RPCError(self._error_message(error))
@staticmethod
def _error_message(error):
msg = error.description
if error.stack_trace:
msg += '\nServer stack trace:\n' + error.stack_trace
return msg
class Decoder(object):
""" Routines for decoding messages and values from
the protocol buffer serialization format """
GUID_LENGTH = 16
_types = Types()
@classmethod
def guid(cls, data):
""" Decode a 16-byte GUID into a string """
return '-'.join(
(hexlify(data[3::-1]), hexlify(data[5:3:-1]),
hexlify(data[7:5:-1]), hexlify(data[8:10]), hexlify(data[10:16])))
@classmethod
def decode(cls, data, typ):
""" Given a python type, and serialized data, decode the value """
if isinstance(typ, MessageType):
return cls.decode_message(data, typ.python_type)
elif isinstance(typ, EnumerationType):
value = cls._decode_value(data, cls._types.sint32_type)
return typ.python_type(value)
elif isinstance(typ, ValueType):
return cls._decode_value(data, typ)
elif isinstance(typ, ClassType):
object_id_typ = cls._types.uint64_type
object_id = cls._decode_value(data, object_id_typ)
return typ.python_type(object_id) if object_id != 0 else None
elif isinstance(typ, ListType):
if data == b'\x00':
return None
msg = cls.decode_message(data, KRPC.List)
return [cls.decode(item, typ.value_type) for item in msg.items]
elif isinstance(typ, DictionaryType):
if data == b'\x00':
return None
msg = cls.decode_message(data, KRPC.Dictionary)
return dict((cls.decode(entry.key, typ.key_type),
cls.decode(entry.value, typ.value_type))
for entry in msg.entries)
elif isinstance(typ, SetType):
if data == b'\x00':
return None
msg = cls.decode_message(data, KRPC.Set)
return set(cls.decode(item, typ.value_type) for item in msg.items)
elif isinstance(typ, TupleType):
if data == b'\x00':
return None
msg = cls.decode_message(data, KRPC.Tuple)
return tuple(
cls.decode(item, value_type)
for item, value_type in zip(msg.items, typ.value_types))
else:
raise EncodingError('Cannot decode type %s' % str(typ))
@classmethod
def decode_message_size(cls, data):
return protobuf_decoder._DecodeVarint(data, 0)[0]
@classmethod
def decode_message(cls, data, typ):
message = typ()
message.ParseFromString(data)
return message
@classmethod
def _decode_value(cls, data, typ):
if typ.protobuf_type.code == KRPC.Type.SINT32:
return _ValueDecoder.decode_sint32(data)
elif typ.protobuf_type.code == KRPC.Type.SINT64:
return _ValueDecoder.decode_sint64(data)
elif typ.protobuf_type.code == KRPC.Type.UINT32:
return _ValueDecoder.decode_uint32(data)
elif typ.protobuf_type.code == KRPC.Type.UINT64:
return _ValueDecoder.decode_uint64(data)
elif typ.protobuf_type.code == KRPC.Type.DOUBLE:
return _ValueDecoder.decode_double(data)
elif typ.protobuf_type.code == KRPC.Type.FLOAT:
return _ValueDecoder.decode_float(data)
elif typ.protobuf_type.code == KRPC.Type.BOOL:
return _ValueDecoder.decode_bool(data)
elif typ.protobuf_type.code == KRPC.Type.STRING:
return _ValueDecoder.decode_string(data)
elif typ.protobuf_type.code == KRPC.Type.BYTES:
return _ValueDecoder.decode_bytes(data)
else:
raise EncodingError('Invalid type')
class Encoder(object):
""" Routines for encoding messages and values in the protocol buffer serialization format """
RPC_HELLO_MESSAGE = b'\x48\x45\x4C\x4C\x4F\x2D\x52\x50\x43\x00\x00\x00'
STREAM_HELLO_MESSAGE = b'\x48\x45\x4C\x4C\x4F\x2D\x53\x54\x52\x45\x41\x4D'
_types = Types()
@classmethod
def client_name(cls, name=None):
""" A client name, truncated/lengthened to fit 32 bytes """
if name is None:
name = ''
else:
name = cls._unicode_truncate(name, 32, 'utf-8')
name = name.encode('utf-8')
return name + (b'\x00' * (32-len(name)))
@classmethod
def _unicode_truncate(cls, string, length, encoding='utf-8'):
""" Shorten a unicode string so that it's encoding uses at
most length bytes. """
encoded = string.encode(encoding=encoding)[:length]
return encoded.decode(encoding, 'ignore')
@classmethod
def encode(cls, x, typ):
""" Encode a message or value of the given protocol buffer type """
if isinstance(typ, MessageType):
return x.SerializeToString()
elif isinstance(typ, ValueType):
return cls._encode_value(x, typ)
elif isinstance(typ, EnumType):
return cls._encode_value(x.value, cls._types.as_type('int32'))
elif isinstance(typ, ClassType):
object_id = x._object_id if x is not None else 0
return cls._encode_value(object_id, cls._types.as_type('uint64'))
elif isinstance(typ, ListType):
msg = cls._types.as_type('KRPC.List').python_type()
msg.items.extend(cls.encode(item, typ.value_type) for item in x)
return msg.SerializeToString()
elif isinstance(typ, DictionaryType):
msg = cls._types.as_type('KRPC.Dictionary').python_type()
entry_type = cls._types.as_type('KRPC.DictionaryEntry')
entries = []
for key, value in sorted(x.items(), key=lambda i: i[0]):
entry = entry_type.python_type()
entry.key = cls.encode(key, typ.key_type)
entry.value = cls.encode(value, typ.value_type)
entries.append(entry)
msg.entries.extend(entries)
return msg.SerializeToString()
elif isinstance(typ, SetType):
msg = cls._types.as_type('KRPC.Set').python_type()
msg.items.extend(cls.encode(item, typ.value_type) for item in x)
return msg.SerializeToString()
elif isinstance(typ, TupleType):
msg = cls._types.as_type('KRPC.Tuple').python_type()
if len(x) != len(typ.value_types):
raise ValueError('Tuple has wrong number of elements. ' +
'Expected %d, got %d.' % (len(typ.value_types), len(x)))
msg.items.extend(cls.encode(item, value_type) for item, value_type in zip(x, typ.value_types))
return msg.SerializeToString()
else:
raise RuntimeError('Cannot encode objects of type ' + str(type(x)))
@classmethod
def encode_delimited(cls, x, typ):
""" Encode a message or value with size information
(for use in a delimited communication stream) """
data = cls.encode(x, typ)
delimiter = protobuf_encoder._VarintBytes(len(data))
return delimiter + data
@classmethod
def _encode_value(cls, value, typ):
return getattr(_ValueEncoder, 'encode_' + typ.protobuf_type)(value)
class Client(object):
"""
A kRPC client, through which all Remote Procedure Calls are made.
Services provided by the server that the client connects to are automatically added.
RPCs can be made using client.ServiceName.ProcedureName(parameter)
"""
def __init__(self, rpc_connection, stream_connection):
self._types = Types()
self._rpc_connection = rpc_connection
self._rpc_connection_lock = threading.Lock()
self._stream_connection = stream_connection
self._request_type = self._types.as_type('KRPC.Request')
self._response_type = self._types.as_type('KRPC.Response')
# Get the services
services = self._invoke(
'KRPC',
'GetServices',
return_type=self._types.as_type('KRPC.Services')).services
# Set up services
for service in services:
setattr(self, snake_case(service.name),
create_service(self, service))
# Set up stream update thread
if stream_connection is not None:
self._stream_thread_stop = threading.Event()
self._stream_thread = threading.Thread(
target=krpc.stream.update_thread,
args=(stream_connection, self._stream_thread_stop))
self._stream_thread.daemon = True
self._stream_thread.start()
else:
self._stream_thread = None
def close(self):
self._rpc_connection.close()
if self._stream_thread is not None:
self._stream_thread_stop.set()
self._stream_thread.join()
def __enter__(self):
return self
def __exit__(self, typ, value, traceback):
self.close()
def add_stream(self, func, *args, **kwargs):
if self._stream_connection is None:
raise RuntimeError('Not connected to stream server')
return krpc.stream.add_stream(self, func, *args, **kwargs)
@contextmanager
def stream(self, func, *args, **kwargs):
""" 'with' support """
s = self.add_stream(func, *args, **kwargs)
try:
yield s
finally:
s.remove()
def _invoke(self,
service,
procedure,
args=[],
kwargs={},
param_names=[],
param_types=[],
return_type=None):
""" Execute an RPC """
# Build the request
request = self._build_request(service, procedure, args, kwargs,
param_names, param_types, return_type)
# Send the request
with self._rpc_connection_lock:
self._send_request(request)
response = self._receive_response()
# Check for an error response
if response.has_error:
raise RPCError(response.error)
# Decode the response and return the (optional) result
result = None
if return_type is not None:
result = Decoder.decode(response.return_value, return_type)
return result
def _build_request(self,
service,
procedure,
args=[],
kwargs={},
param_names=[],
param_types=[],
return_type=None):
""" Build a KRPC.Request object """
def encode_argument(i, value):
typ = param_types[i]
if type(value) != typ.python_type:
# Try coercing to the correct type
try:
value = self._types.coerce_to(value, typ)
except ValueError:
raise TypeError('%s.%s() argument %d must be a %s, got a %s' % \
(service, procedure, i, typ.python_type, type(value)))
return Encoder.encode(value, typ)
if len(args) > len(param_types):
raise TypeError('%s.%s() takes exactly %d arguments (%d given)' % \
(service, procedure, len(param_types), len(args)))
arguments = []
nargs = len(args)
for i, param in enumerate(param_names):
add = False
if i < nargs and not isinstance(args[i], DefaultArgument):
arg = args[i]
add = True
elif param in kwargs:
arg = kwargs[param]
add = True
if add:
argument = krpc.schema.KRPC.Argument()
argument.position = i
argument.value = encode_argument(i, arg)
arguments.append(argument)
# Build the request object
request = krpc.schema.KRPC.Request()
request.service = service
request.procedure = procedure
request.arguments.extend(arguments)
return request
def _send_request(self, request):
""" Send a KRPC.Request object to the server """
data = Encoder.encode_delimited(request, self._request_type)
self._rpc_connection.send(data)
def _receive_response(self):
""" Receive data from the server and decode it into a KRPC.Response object """
# Read the size and position of the response message
data = b''
while True:
try:
data += self._rpc_connection.partial_receive(1)
size, position = Decoder.decode_size_and_position(data)
break
except IndexError:
pass
# Read and decode the response message
data = self._rpc_connection.receive(size)
return Decoder.decode(data, self._response_type)
class TestEncodeDecode(unittest.TestCase):
types = Types()
def _run_test_encode_value(self, typ, cases):
for decoded, encoded in cases:
data = Encoder.encode(decoded, typ)
self.assertEqual(encoded, hexlify(data))
def _run_test_decode_value(self, typ, cases):
for decoded, encoded in cases:
value = Decoder.decode(unhexlify(encoded), typ)
if typ.python_type == float:
self.assertEqual(str(decoded)[0:8], str(value)[0:8])
else:
self.assertEqual(decoded, value)
def test_double(self):
cases = [(0.0, '0000000000000000'), (-1.0, '000000000000f0bf'),
(3.14159265359, 'ea2e4454fb210940'),
(float('inf'), '000000000000f07f'),
(-float('inf'), '000000000000f0ff'),
(float('nan'), '000000000000f87f')]
self._run_test_encode_value(self.types.double_type, cases)
self._run_test_decode_value(self.types.double_type, cases)
def test_float(self):
cases = [(3.14159265359, 'db0f4940'), (-1.0, '000080bf'),
(0.0, '00000000'), (float('inf'), '0000807f'),
(-float('inf'), '000080ff'), (float('nan'), '0000c07f')]
self._run_test_encode_value(self.types.float_type, cases)
self._run_test_decode_value(self.types.float_type, cases)
def test_sint32(self):
cases = [(0, '00'), (1, '02'), (42, '54'), (300, 'd804'), (-33, '41'),
(2147483647, 'feffffff0f'), (-2147483648, 'ffffffff0f')]
self._run_test_encode_value(self.types.sint32_type, cases)
self._run_test_decode_value(self.types.sint32_type, cases)
def test_sint64(self):
cases = [(0, '00'), (1, '02'), (42, '54'), (300, 'd804'),
(1234567890000L, 'a091d89fee47'), (-33, '41')]
self._run_test_encode_value(self.types.sint64_type, cases)
self._run_test_decode_value(self.types.sint64_type, cases)
def test_uint32(self):
cases = [(0, '00'), (1, '01'), (42, '2a'), (300, 'ac02'),
(sys.maxint, 'ffffffffffffffff7f')]
self._run_test_encode_value(self.types.uint32_type, cases)
self._run_test_decode_value(self.types.uint32_type, cases)
self.assertRaises(EncodingError, Encoder.encode, -1,
self.types.uint32_type)
self.assertRaises(EncodingError, Encoder.encode, -849,
self.types.uint32_type)
def test_uint64(self):
cases = [(0, '00'), (1, '01'), (42, '2a'), (300, 'ac02'),
(1234567890000L, 'd088ec8ff723')]
self._run_test_encode_value(self.types.uint64_type, cases)
self._run_test_decode_value(self.types.uint64_type, cases)
self.assertRaises(EncodingError, Encoder.encode, -1,
self.types.uint64_type)
self.assertRaises(EncodingError, Encoder.encode, -849,
self.types.uint64_type)
def test_bool(self):
cases = [(True, '01'), (False, '00')]
self._run_test_encode_value(self.types.bool_type, cases)
self._run_test_decode_value(self.types.bool_type, cases)
def test_string(self):
cases = [('', '00'), ('testing', '0774657374696e67'),
('One small step for Kerbal-kind!',
'1f4f6e6520736d616c6c207374657020' +
'666f72204b657262616c2d6b696e6421'),
(b'\xe2\x84\xa2'.decode('utf-8'), '03e284a2'),
(b'Mystery Goo\xe2\x84\xa2 Containment Unit'.decode('utf-8'),
'1f4d79737465727920476f6fe284a220' +
'436f6e7461696e6d656e7420556e6974')]
self._run_test_encode_value(self.types.string_type, cases)
self._run_test_decode_value(self.types.string_type, cases)
def test_bytes(self):
cases = [(b'', '00'), (b'\xba\xda\x55', '03bada55'),
(b'\xde\xad\xbe\xef', '04deadbeef')]
self._run_test_encode_value(self.types.bytes_type, cases)
self._run_test_decode_value(self.types.bytes_type, cases)
def test_tuple(self):
cases = [((1, ), '0a0101')]
self._run_test_encode_value(
self.types.tuple_type(self.types.uint32_type), cases)
self._run_test_decode_value(
self.types.tuple_type(self.types.uint32_type), cases)
cases = [((1, 'jeb', False), '0a01010a04036a65620a0100')]
typ = self.types.tuple_type(self.types.uint32_type,
self.types.string_type,
self.types.bool_type)
self._run_test_encode_value(typ, cases)
self._run_test_decode_value(typ, cases)
def test_list(self):
cases = [([], ''), ([1], '0a0101'),
([1, 2, 3, 4], '0a01010a01020a01030a0104')]
typ = self.types.list_type(self.types.uint32_type)
self._run_test_encode_value(typ, cases)
self._run_test_decode_value(typ, cases)
def test_set(self):
cases = [(set(), ''), (set([1]), '0a0101'),
(set([1, 2, 3, 4]), '0a01010a01020a01030a0104')]
typ = self.types.set_type(self.types.uint32_type)
self._run_test_encode_value(typ, cases)
self._run_test_decode_value(typ, cases)
def test_dictionary(self):
cases = [({}, ''), ({
'': 0
}, '0a060a0100120100'),
({
'foo': 42,
'bar': 365,
'baz': 3
}, '0a0a0a04036261721202ed020a090a0403' +
'62617a1201030a090a0403666f6f12012a')]
typ = self.types.dictionary_type(self.types.string_type,
self.types.uint32_type)
self._run_test_encode_value(typ, cases)
self._run_test_decode_value(typ, cases)
class Encoder(object):
""" Routines for encoding messages and values in
the protocol buffer serialization format """
_types = Types()
@classmethod
def encode(cls, x, typ):
""" Encode a message or value of the given protocol buffer type """
if isinstance(typ, MessageType):
return x.SerializeToString()
elif isinstance(typ, ValueType):
return cls._encode_value(x, typ)
elif isinstance(typ, EnumerationType):
return cls._encode_value(x.value, cls._types.sint32_type)
elif isinstance(typ, ClassType):
object_id = x._object_id if x is not None else 0
return cls._encode_value(object_id, cls._types.uint64_type)
elif isinstance(typ, ListType):
msg = KRPC.List()
msg.items.extend(cls.encode(item, typ.value_type) for item in x)
return msg.SerializeToString()
elif isinstance(typ, DictionaryType):
msg = KRPC.Dictionary()
entries = []
for key, value in sorted(list(x.items()), key=lambda i: i[0]):
entry = KRPC.DictionaryEntry()
entry.key = cls.encode(key, typ.key_type)
entry.value = cls.encode(value, typ.value_type)
entries.append(entry)
msg.entries.extend(entries)
return msg.SerializeToString()
elif isinstance(typ, SetType):
msg = KRPC.Set()
msg.items.extend(cls.encode(item, typ.value_type) for item in x)
return msg.SerializeToString()
elif isinstance(typ, TupleType):
msg = KRPC.Tuple()
if len(x) != len(typ.value_types):
raise EncodingError('Tuple has wrong number of elements. ' +
'Expected %d, got %d.' %
(len(typ.value_types), len(x)))
msg.items.extend(
cls.encode(item, value_type)
for item, value_type in zip(x, typ.value_types))
return msg.SerializeToString()
else:
raise EncodingError('Cannot encode objects of type ' +
str(type(x)))
@classmethod
def encode_message_with_size(cls, message):
""" Encode a protobuf message, prepended with its size """
data = message.SerializeToString()
size = protobuf_encoder._VarintBytes(len(data))
return size + data
@classmethod
def _encode_value(cls, value, typ):
if typ.protobuf_type.code == KRPC.Type.SINT32:
return _ValueEncoder.encode_sint32(value)
elif typ.protobuf_type.code == KRPC.Type.SINT64:
return _ValueEncoder.encode_sint64(value)
elif typ.protobuf_type.code == KRPC.Type.UINT32:
return _ValueEncoder.encode_uint32(value)
elif typ.protobuf_type.code == KRPC.Type.UINT64:
return _ValueEncoder.encode_uint64(value)
elif typ.protobuf_type.code == KRPC.Type.DOUBLE:
return _ValueEncoder.encode_double(value)
elif typ.protobuf_type.code == KRPC.Type.FLOAT:
return _ValueEncoder.encode_float(value)
elif typ.protobuf_type.code == KRPC.Type.BOOL:
return _ValueEncoder.encode_bool(value)
elif typ.protobuf_type.code == KRPC.Type.STRING:
return _ValueEncoder.encode_string(value)
elif typ.protobuf_type.code == KRPC.Type.BYTES:
return _ValueEncoder.encode_bytes(value)
else:
raise EncodingError('Invalid type')
def test_coerce_to(self):
types = Types()
cases = [
(42.0, 42, 'double'),
(42.0, 42, 'float'),
(42, 42.0, 'int32'),
(42, 42L, 'int32'),
(42L, 42.0, 'int64'),
(42L, 42, 'int64'),
(42, 42.0, 'uint32'),
(42, 42L, 'uint32'),
(42L, 42.0, 'uint64'),
(42L, 42, 'uint64'),
(list(), tuple(), 'List(string)'),
((0, 1, 2), [0, 1, 2], 'Tuple(int32,int32,int32)'),
([0, 1, 2], (0, 1, 2), 'List(int32)'),
]
for expected, value, typ in cases:
coerced_value = types.coerce_to(value, types.as_type(typ))
self.assertEqual(expected, coerced_value)
self.assertEqual(type(expected), type(coerced_value))
strings = [
u'foo',
u'\xe2\x84\xa2',
u'Mystery Goo\xe2\x84\xa2 Containment Unit'
]
for string in strings:
self.assertEqual(string, types.coerce_to(string, types.as_type('string')))
self.assertEqual(['foo', 'bar'], types.coerce_to(['foo', 'bar'], types.as_type('List(string)')))
self.assertRaises(ValueError, types.coerce_to, None, types.as_type('float'))
self.assertRaises(ValueError, types.coerce_to, '', types.as_type('float'))
self.assertRaises(ValueError, types.coerce_to, True, types.as_type('float'))
self.assertRaises(ValueError, types.coerce_to, list(), types.as_type('Tuple(int32)'))
self.assertRaises(ValueError, types.coerce_to, ['foo', 2], types.as_type('Tuple(string)'))
self.assertRaises(ValueError, types.coerce_to, [1], types.as_type('Tuple(string)'))
class TestEncoder(unittest.TestCase):
types = Types()
def test_rpc_hello_message(self):
message = Encoder.RPC_HELLO_MESSAGE
self.assertEqual(12, len(message))
self.assertEqual('48454c4c4f2d525043000000', hexlify(message))
def test_stream_hello_message(self):
message = Encoder.STREAM_HELLO_MESSAGE
self.assertEqual(12, len(message))
self.assertEqual('48454c4c4f2d53545245414d', hexlify(message))
def test_client_name(self):
message = Encoder.client_name('foo')
self.assertEqual(32, len(message))
self.assertEqual('666f6f' + '00' * 29, hexlify(message))
def test_empty_client_name(self):
message = Encoder.client_name()
self.assertEqual(32, len(message))
self.assertEqual('00' * 32, hexlify(message))
def test_long_client_name(self):
message = Encoder.client_name('a' * 33)
self.assertEqual(32, len(message))
self.assertEqual('61' * 32, hexlify(message))
def test_encode_message(self):
request = KRPC.Request()
request.service = 'ServiceName'
request.procedure = 'ProcedureName'
data = Encoder.encode(request, self.types.as_type('KRPC.Request'))
expected = '0a0b536572766963654e616d65120d50726f6365647572654e616d65'
self.assertEqual(expected, hexlify(data))
def test_encode_value(self):
data = Encoder.encode(300, self.types.as_type('int32'))
self.assertEqual('ac02', hexlify(data))
def test_encode_unicode_string(self):
data = Encoder.encode(
b'\xe2\x84\xa2'.decode('utf-8'), self.types.as_type('string'))
self.assertEqual('03e284a2', hexlify(data))
def test_encode_message_delimited(self):
request = KRPC.Request()
request.service = 'ServiceName'
request.procedure = 'ProcedureName'
data = Encoder.encode_delimited(
request, self.types.as_type('KRPC.Request'))
expected = '1c' + \
'0a0b536572766963654e616d6512' + \
'0d50726f6365647572654e616d65'
self.assertEqual(expected, hexlify(data))
def test_encode_value_delimited(self):
data = Encoder.encode_delimited(300, self.types.as_type('int32'))
self.assertEqual('02' + 'ac02', hexlify(data))
def test_encode_class(self):
typ = self.types.as_type('Class(ServiceName.ClassName)')
class_type = typ.python_type
self.assertTrue(issubclass(class_type, ClassBase))
value = class_type(300)
self.assertEqual(300, value._object_id)
data = Encoder.encode(value, typ)
self.assertEqual('ac02', hexlify(data))
def test_encode_class_none(self):
typ = self.types.as_type('Class(ServiceName.ClassName)')
value = None
data = Encoder.encode(value, typ)
self.assertEqual('00', hexlify(data))
def test_encode_tuple_wrong_arity(self):
typ = self.types.as_type('Tuple(int32,int32,int32)')
value = (0, 1)
self.assertRaises(ValueError, Encoder.encode, value, typ)
def test_coerce_to(self):
types = Types()
cases = [
(42.0, 42, 'double'),
(42.0, 42, 'float'),
(42, 42.0, 'int32'),
(42, 42L, 'int32'),
(42L, 42.0, 'int64'),
(42L, 42, 'int64'),
(42, 42.0, 'uint32'),
(42, 42L, 'uint32'),
(42L, 42.0, 'uint64'),
(42L, 42, 'uint64'),
(list(), tuple(), 'List(string)'),
((0, 1, 2), [0, 1, 2], 'Tuple(int32,int32,int32)'),
([0, 1, 2], (0, 1, 2), 'List(int32)'),
]
for expected, value, typ in cases:
coerced_value = types.coerce_to(value, types.as_type(typ))
self.assertEqual(expected, coerced_value)
self.assertEqual(type(expected), type(coerced_value))
strings = [
u'foo', u'\xe2\x84\xa2',
u'Mystery Goo\xe2\x84\xa2 Containment Unit'
]
for string in strings:
self.assertEqual(string,
types.coerce_to(string, types.as_type('string')))
self.assertEqual(['foo', 'bar'],
types.coerce_to(['foo', 'bar'],
types.as_type('List(string)')))
self.assertRaises(ValueError, types.coerce_to, None,
types.as_type('float'))
self.assertRaises(ValueError, types.coerce_to, '',
types.as_type('float'))
self.assertRaises(ValueError, types.coerce_to, True,
types.as_type('float'))
self.assertRaises(ValueError, types.coerce_to, list(),
types.as_type('Tuple(int32)'))
self.assertRaises(ValueError, types.coerce_to, ['foo', 2],
types.as_type('Tuple(string)'))
self.assertRaises(ValueError, types.coerce_to, [1],
types.as_type('Tuple(string)'))
class Client(object):
"""
A kRPC client, through which all Remote Procedure Calls are made.
Services provided by the server that the client connects
to are automatically added. RPCs can be made using
client.ServiceName.ProcedureName(parameter)
"""
def __init__(self, rpc_connection, stream_connection):
self._types = Types()
self._rpc_connection = rpc_connection
self._rpc_connection_lock = threading.Lock()
self._stream_connection = stream_connection
self._stream_cache = {}
self._stream_cache_lock = threading.Lock()
self._request_type = self._types.as_type('KRPC.Request')
self._response_type = self._types.as_type('KRPC.Response')
# Get the services
services = self._invoke('KRPC', 'GetServices', [], [], [],
self._types.as_type('KRPC.Services')).services
# Set up services
for service in services:
setattr(self, snake_case(service.name),
create_service(self, service))
# Set up stream update thread
if stream_connection is not None:
self._stream_thread_stop = threading.Event()
self._stream_thread = threading.Thread(
target=krpc.stream.update_thread,
args=(stream_connection, self._stream_thread_stop,
self._stream_cache, self._stream_cache_lock))
self._stream_thread.daemon = True
self._stream_thread.start()
else:
self._stream_thread = None
def close(self):
self._rpc_connection.close()
if self._stream_thread is not None:
self._stream_thread_stop.set()
self._stream_thread.join()
def __enter__(self):
return self
def __exit__(self, typ, value, traceback):
self.close()
def add_stream(self, func, *args, **kwargs):
if self._stream_connection is None:
raise RuntimeError('Not connected to stream server')
return krpc.stream.add_stream(self, func, *args, **kwargs)
@contextmanager
def stream(self, func, *args, **kwargs):
""" 'with' support """
stream = self.add_stream(func, *args, **kwargs)
try:
yield stream
finally:
stream.remove()
def _invoke(self, service, procedure, args, param_names, param_types,
return_type):
""" Execute an RPC """
# Build the request
request = self._build_request(service, procedure, args, param_names,
param_types, return_type)
# Send the request
with self._rpc_connection_lock:
self._send_request(request)
response = self._receive_response()
# Check for an error response
if response.has_error:
raise RPCError(response.error)
# Decode the response and return the (optional) result
result = None
if return_type is not None:
result = Decoder.decode(response.return_value, return_type)
return result
# pylint: disable=unused-argument
def _build_request(self, service, procedure, args, param_names,
param_types, return_type):
""" Build a KRPC.Request object """
request = krpc.schema.KRPC.Request(service=service,
procedure=procedure)
for i, (value, typ) in enumerate(itertools.izip(args, param_types)):
if isinstance(value, DefaultArgument):
continue
if not isinstance(value, typ.python_type):
try:
value = self._types.coerce_to(value, typ)
except ValueError:
raise TypeError(
'%s.%s() argument %d must be a %s, got a %s' %
(service, procedure, i, typ.python_type, type(value)))
request.arguments.add(position=i, value=Encoder.encode(value, typ))
return request
def _send_request(self, request):
""" Send a KRPC.Request object to the server """
data = Encoder.encode_delimited(request, self._request_type)
self._rpc_connection.send(data)
def _receive_response(self):
""" Receive data from the server and
decode it into a KRPC.Response object """
# Read the size and position of the response message
data = b''
while True:
try:
data += self._rpc_connection.partial_receive(1)
size, _ = Decoder.decode_size_and_position(data)
break
except IndexError:
pass
# Read and decode the response message
data = self._rpc_connection.receive(size)
return Decoder.decode(data, self._response_type)
class Client(object):
"""
A kRPC client, through which all Remote Procedure Calls are made.
Services provided by the server that the client connects to are automatically added.
RPCs can be made using client.ServiceName.ProcedureName(parameter)
"""
def __init__(self, rpc_connection, stream_connection):
self._types = Types()
self._rpc_connection = rpc_connection
self._rpc_connection_lock = threading.Lock()
self._stream_connection = stream_connection
self._stream_cache = {}
self._stream_cache_lock = threading.Lock()
self._request_type = self._types.as_type('KRPC.Request')
self._response_type = self._types.as_type('KRPC.Response')
# Get the services
services = self._invoke('KRPC', 'GetServices', [], [], [], self._types.as_type('KRPC.Services')).services
# Set up services
for service in services:
setattr(self, snake_case(service.name), create_service(self, service))
# Set up stream update thread
if stream_connection is not None:
self._stream_thread_stop = threading.Event()
self._stream_thread = threading.Thread(target=krpc.stream.update_thread,
args=(stream_connection, self._stream_thread_stop,
self._stream_cache, self._stream_cache_lock))
self._stream_thread.daemon = True
self._stream_thread.start()
else:
self._stream_thread = None
def close(self):
self._rpc_connection.close()
if self._stream_thread is not None:
self._stream_thread_stop.set()
self._stream_thread.join()
def __enter__(self):
return self
def __exit__(self, typ, value, traceback):
self.close()
def add_stream(self, func, *args, **kwargs):
if self._stream_connection is None:
raise RuntimeError('Not connected to stream server')
return krpc.stream.add_stream(self, func, *args, **kwargs)
@contextmanager
def stream(self, func, *args, **kwargs):
""" 'with' support """
stream = self.add_stream(func, *args, **kwargs)
try:
yield stream
finally:
stream.remove()
def _invoke(self, service, procedure, args, param_names, param_types, return_type):
""" Execute an RPC """
# Build the request
request = self._build_request(service, procedure, args, param_names, param_types, return_type)
# Send the request
with self._rpc_connection_lock:
self._send_request(request)
response = self._receive_response()
# Check for an error response
if response.has_error:
raise RPCError(response.error)
# Decode the response and return the (optional) result
result = None
if return_type is not None:
result = Decoder.decode(response.return_value, return_type)
return result
def _build_request(self, service, procedure, args,
param_names, param_types, return_type): #pylint: disable=unused-argument
""" Build a KRPC.Request object """
request = krpc.schema.KRPC.Request(service=service, procedure=procedure)
for i, (value, typ) in enumerate(itertools.izip(args, param_types)):
if isinstance(value, DefaultArgument):
continue
if not isinstance(value, typ.python_type):
try:
value = self._types.coerce_to(value, typ)
except ValueError:
raise TypeError('%s.%s() argument %d must be a %s, got a %s' % \
(service, procedure, i, typ.python_type, type(value)))
request.arguments.add(position=i, value=Encoder.encode(value, typ))
return request
def _send_request(self, request):
""" Send a KRPC.Request object to the server """
data = Encoder.encode_delimited(request, self._request_type)
self._rpc_connection.send(data)
def _receive_response(self):
""" Receive data from the server and decode it into a KRPC.Response object """
# Read the size and position of the response message
data = b''
while True:
try:
data += self._rpc_connection.partial_receive(1)
size, _ = Decoder.decode_size_and_position(data)
break
except IndexError:
pass
# Read and decode the response message
data = self._rpc_connection.receive(size)
return Decoder.decode(data, self._response_type)
