def _unpack_arg(arg_types, kwarg_types, arg) -> _Arguments:
"""Unpacks 'arg' into an argument list based on types."""
args = []
for idx, expected_type in enumerate(arg_types):
element_value = arg[idx]
actual_type = type_conversions.infer_type(element_value)
if not expected_type.is_assignable_from(actual_type):
raise TypeError(
'Expected element at position {} to be of type {}, found {}.'.
format(idx, expected_type, actual_type))
if isinstance(element_value, structure.Struct):
element_value = type_conversions.type_to_py_container(
element_value, expected_type)
args.append(element_value)
kwargs = {}
for name, expected_type in kwarg_types.items():
element_value = getattr(arg, name)
actual_type = type_conversions.infer_type(element_value)
if not expected_type.is_assignable_from(actual_type):
raise TypeError(
'Expected element named {} to be of type {}, found {}.'.format(
name, expected_type, actual_type))
if type_analysis.is_struct_with_py_container(element_value,
expected_type):
element_value = type_conversions.type_to_py_container(
element_value, expected_type)
kwargs[name] = element_value
return args, kwargs
def _ensure_arg_type(parameter_type, arg) -> _Arguments:
"""Ensures that `arg` matches `parameter_type` before returning it."""
arg_type = type_conversions.infer_type(arg)
if not parameter_type.is_assignable_from(arg_type):
raise TypeError('Expected an argument of type {}, found {}.'.format(
parameter_type, arg_type))
if type_analysis.is_struct_with_py_container(arg, parameter_type):
arg = type_conversions.type_to_py_container(arg, parameter_type)
return [arg], {}
def _call(fn, parameter_type, arg):
arg_type = type_conversions.infer_type(arg)
if not parameter_type.is_assignable_from(arg_type):
raise TypeError(
'Expected an argument of type {}, found {}.'.format(
parameter_type, arg_type))
if type_analysis.is_struct_with_py_container(
arg, parameter_type):
arg = type_conversions.type_to_py_container(
arg, parameter_type)
return fn(arg)
def _unpack_and_call(fn, arg_types, kwarg_types, arg):
"""An interceptor function that unpacks 'arg' before calling `fn`.
The function verifies the actual parameters before it forwards the
call as a last-minute check.
Args:
fn: The function or defun to invoke.
arg_types: The list of positional argument types (guaranteed to all be
instances of computation_types.Types).
kwarg_types: The dictionary of keyword argument types (guaranteed to
all be instances of computation_types.Types).
arg: The argument to unpack.
Returns:
The result of invoking `fn` on the unpacked arguments.
Raises:
TypeError: if types don't match.
"""
py_typecheck.check_type(arg,
(structure.Struct, value_base.Value))
args = []
for idx, expected_type in enumerate(arg_types):
element_value = arg[idx]
actual_type = type_conversions.infer_type(element_value)
if not expected_type.is_assignable_from(actual_type):
raise TypeError(
'Expected element at position {} to be of type {}, found {}.'
.format(idx, expected_type, actual_type))
if isinstance(element_value, structure.Struct):
element_value = type_conversions.type_to_py_container(
element_value, expected_type)
args.append(element_value)
kwargs = {}
for name, expected_type in kwarg_types.items():
element_value = getattr(arg, name)
actual_type = type_conversions.infer_type(element_value)
if not expected_type.is_assignable_from(actual_type):
raise TypeError(
'Expected element named {} to be of type {}, found {}.'
.format(name, expected_type, actual_type))
if type_analysis.is_struct_with_py_container(
element_value, expected_type):
element_value = type_conversions.type_to_py_container(
element_value, expected_type)
kwargs[name] = element_value
return fn(*args, **kwargs)
def test_returns_false_with_named_tuple_type_spec(self):
value = structure.Struct([('a', 0.0)])
type_spec = computation_types.StructType([('a', tf.float32)])
self.assertFalse(
type_analysis.is_struct_with_py_container(value, type_spec))
def test_returns_false_with_none_value(self):
value = None
type_spec = computation_types.StructWithPythonType([('a', tf.float32)],
dict)
self.assertFalse(
type_analysis.is_struct_with_py_container(value, type_spec))
def test_returns_true(self):
value = structure.Struct([('a', 0.0)])
type_spec = computation_types.StructWithPythonType([('a', tf.float32)],
dict)
self.assertTrue(
type_analysis.is_struct_with_py_container(value, type_spec))
def test_returns_true(self):
value = anonymous_tuple.AnonymousTuple([('a', 0.0)])
type_spec = computation_types.NamedTupleTypeWithPyContainerType(
[('a', tf.float32)], dict)
self.assertTrue(type_analysis.is_struct_with_py_container(value, type_spec))
