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() } }