def _arg_to_proto()

in tensorflow_quantum/core/serialize/op_serializer.py [0:0]

• 44 lines of code
• 18 McCabe index (conditional complexity)

def _arg_to_proto(value, *, arg_function_language, out=None):
    """Writes an argument value into an Arg proto.
    Args:
        value: The value to encode.
        arg_function_language: The language to use when encoding functions. If
            this is set to None, it will be set to the minimal language
            necessary to support the features that were actually used.
        out: The proto to write the result into. Defaults to a new instance.
    Returns:
        The proto that was written into as well as the `arg_function_language`
        that was used.
    """

    if arg_function_language not in SUPPORTED_FUNCTIONS_FOR_LANGUAGE:
        raise ValueError(f'Unrecognized arg_function_language: '
                         f'{arg_function_language!r}')
    supported = SUPPORTED_FUNCTIONS_FOR_LANGUAGE[arg_function_language]

    msg = program_pb2.Arg() if out is None else out

    def check_support(func_type: str) -> str:
        if func_type not in supported:
            lang = (repr(arg_function_language)
                    if arg_function_language is not None else '[any]')
            raise ValueError(f'Function type {func_type!r} not supported by '
                             f'arg_function_language {lang}')
        return func_type

    if isinstance(value, (float, int, sympy.Integer, sympy.Float,
                          sympy.Rational, sympy.NumberSymbol)):
        msg.arg_value.float_value = float(value)
    elif isinstance(value, str):
        msg.arg_value.string_value = value
    elif (isinstance(value, (list, tuple, np.ndarray)) and
          all(isinstance(x, (bool, np.bool_)) for x in value)):
        # Some protobuf / numpy combinations do not support np.bool_, so cast.
        msg.arg_value.bool_values.values.extend([bool(x) for x in value])
    elif isinstance(value, sympy.Symbol):
        msg.symbol = str(value.free_symbols.pop())
    elif isinstance(value, sympy.Add):
        msg.func.type = check_support('add')
        for arg in value.args:
            _arg_to_proto(arg,
                          arg_function_language=arg_function_language,
                          out=msg.func.args.add())
    elif isinstance(value, sympy.Mul):
        msg.func.type = check_support('mul')
        for arg in value.args:
            _arg_to_proto(arg,
                          arg_function_language=arg_function_language,
                          out=msg.func.args.add())
    elif isinstance(value, sympy.Pow):
        msg.func.type = check_support('pow')
        for arg in value.args:
            _arg_to_proto(arg,
                          arg_function_language=arg_function_language,
                          out=msg.func.args.add())
    else:
        raise ValueError(f'Unrecognized arg type: {type(value)}')

    return msg