def __init__(self, func, module):
self._func = func
self._module = module
self._input_type = func.input_type
self._output_type = func.output_type
self._pb_input_type = getattr(
module, func.input_type.__name__) if _is_namedtuple(
self._input_type) else None
self._pb_output_type = getattr(
module, func.output_type.__name__) if _is_namedtuple(
self._output_type) else None
def _gen_service(model, name='Model'):
'''Returns a protobuf service definition string'''
def _type_name(t: type):
"we use empty as placeholders for raw types in protobuf"
return t.__name__ if not is_raw_type(t) else Empty.__name__
rpc_comps = (
(n, _type_name(f.input_type), _type_name(f.output_type))
for n, f in model.methods.items()
if _is_namedtuple(f.input_type) or _is_namedtuple(f.output_type))
rpc_defs = '\n'.join(_gen_rpc(*comps) for comps in rpc_comps)
return _service_template.format(name=name, service_def=rpc_defs)
def _proto_iter(nt):
'''Recursively yields all types contained within the NamedTuple relevant to protobuf gen'''
if _is_namedtuple(nt):
yield nt
for t in nt._field_types.values():
if _is_namedtuple(t):
yield from _proto_iter(t)
elif issubclass(t, Enum):
yield t
elif issubclass(t, List) or issubclass(t, Dict):
for tt in t.__args__:
yield from _proto_iter(tt)
def _get_pb_value(wrapped_type, pb_value):
'''Recursively traverses the protobuf message to ensure nested messages become NamedTuples'''
if _is_namedtuple(wrapped_type):
return _unpack_pb_msg(wrapped_type, pb_value)
elif issubclass(wrapped_type, Dict):
_, val_type = wrapped_type.__args__
if _is_namedtuple(val_type):
return {k: _unpack_pb_msg(val_type, v) for k, v in pb_value.items()}
elif issubclass(wrapped_type, List):
list_type = wrapped_type.__args__[0]
if _is_namedtuple(list_type):
return [_unpack_pb_msg(list_type, v) for v in pb_value]
return pb_value
def _set_pb_value(wrapped_type, value, module):
'''Recursively traverses the NamedTuple instance to ensure nested NamedTuples become protobuf messages'''
if _is_namedtuple(wrapped_type):
return _pack_pb_msg(value, module)
elif issubclass(wrapped_type, Dict):
_, val_type = wrapped_type.__args__
if _is_namedtuple(val_type):
return {k: _pack_pb_msg(v, module) for k, v in value.items()}
elif issubclass(wrapped_type, List):
list_type = wrapped_type.__args__[0]
if _is_namedtuple(list_type):
return [_pack_pb_msg(v, module) for v in value]
return value
def _type2proto(t):
'''Returns a string corresponding to the protobuf type'''
if t in _type_lookup:
return _type_lookup[t]
elif _is_namedtuple(t) or issubclass(t, Enum):
return t.__name__
else:
raise AcumosError("Unknown protobuf mapping for type {}".format(t))
def wrapped_save(pickler, obj, save_persistent_id=True):
'''Hook that intercepts objects about to be saved'''
_catch_object(obj)
if _is_namedtuple(obj) and obj is not Empty:
_save_namedtuple(pickler, obj)
else:
pickler_save(pickler, obj, save_persistent_id)
def _types_equal(t1, t2):
'''Returns True if t1 and t2 types are equal. Can't override __eq__ on NamedTuple unfortunately.'''
if _is_namedtuple(t1) and _is_namedtuple(t2):
names_match = t1.__name__ == t2.__name__
ft1, ft2 = namedtuple_field_types(t1), namedtuple_field_types(t2)
keys_match = ft1.keys() == ft2.keys()
values_match = all(
_types_equal(v1, v2) for v1, v2 in zip(ft1.values(), ft2.values()))
return names_match and keys_match and values_match
if issubclass(t1, Enum) and issubclass(t2, Enum):
names_match = t1.__name__ == t2.__name__
enums_match = [(e.name, e.value) for e in t1] == [(e.name, e.value)
for e in t2]
return names_match and enums_match
else:
return t1 == t2
def _proto_iter(nt):
'''Recursively yields all types contained within the NamedTuple relevant to protobuf gen'''
if is_raw_type(nt):
# Empty is used here as a placeholder for raw types
yield Empty
return
if _is_namedtuple(nt):
yield nt
for t in nt.__annotations__.values():
inspected = inspect_type(t)
if _is_namedtuple(t):
yield from _proto_iter(t)
elif issubclass(inspected.origin, Enum):
yield t
elif issubclass(inspected.origin, List) or issubclass(
inspected.origin, Dict):
for tt in inspected.args:
yield from _proto_iter(tt)
def model2proto(model, package_name):
'''Converts a Model object to a protobuf schema string'''
all_types = (iterchain(
*(iterchain(_proto_iter(f.input_type), _proto_iter(f.output_type))
for f in model.methods.values())))
unique_types = _require_unique(all_types)
type_names = set(t.__name__ for t in unique_types)
msg_defs = tuple(
_nt2proto(t, type_names) if _is_namedtuple(t) else _enum2proto(t)
for t in unique_types)
service_def = _gen_service(model)
package_def = _package_template.format(name=package_name)
defs = (_PROTO_SYNTAX, package_def, service_def) + msg_defs
return '\n'.join(defs)
def _field2proto(name, type_, index, type_names, rjust=None):
'''Returns a protobuf schema field str from a NamedTuple field'''
string = None
inspected = inspect_type(type_)
if type_ in _type_lookup:
string = "{} {} = {};".format(_type2proto(type_), name, index)
elif _is_namedtuple(type_) or issubclass(inspected.origin, Enum):
tn = type_.__name__
if tn not in type_names:
raise AcumosError(
"Could not build protobuf field using unknown custom type {}".
format(tn))
string = "{} {} = {};".format(tn, name, index)
elif issubclass(inspected.origin, List):
inner = inspected.args[0]
if _is_container(inner):
raise NestedTypeError(
"Nested container {} is not yet supported; try using NamedTuple instead"
.format(type_))
string = "repeated {}".format(
_field2proto(name, inner, index, type_names, 0))
elif issubclass(inspected.origin, Dict):
k, v = inspected.args
if any(map(_is_container, (k, v))):
raise NestedTypeError(
"Nested container {} is not yet supported; try using NamedTuple instead"
.format(type_))
string = "map<{}, {}> {} = {};".format(_type2proto(k), _type2proto(v),
name, index)
if string is None:
raise AcumosError(
"Could not build protobuf field due to unsupported type {}".format(
type_))
if rjust is None:
rjust = len(string) + 2
return string.rjust(rjust, ' ')
def create_model_meta(model, name, requirements, encoding='protobuf'):
'''Returns a model metadata dictionary'''
return {
'schema': _SCHEMA,
'runtime': _create_runtime(requirements, encoding),
'name': name,
'methods': {
name: {
'input': {
'name':
f.input_type.__name__,
'media_type': [
_MEDIA_TYPES[encoding if _is_namedtuple(f.input_type
) else f.
input_type.__supertype__._raw_type]
],
'metadata': {} if _is_namedtuple(f.input_type) else
f.input_type.__supertype__._metadata,
'description':
'' if _is_namedtuple(
f.input_type) else f.input_type.__supertype__._doc
},
'output': {
'name':
f.output_type.__name__,
'media_type': [
_MEDIA_TYPES[encoding if _is_namedtuple(f.output_type
) else f.
output_type.__supertype__._raw_type]
],
'metadata': {} if _is_namedtuple(f.output_type) else
f.output_type.__supertype__._metadata,
'description':
'' if _is_namedtuple(
f.output_type) else f.output_type.__supertype__._doc
},
'description': f.description
}
for name, f in model.methods.items()
}
}
