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, (anonymous_tuple.AnonymousTuple, value_base.Value))
args = []
for idx, expected_type in enumerate(arg_types):
element_value = arg[idx]
actual_type = type_utils.infer_type(element_value)
if not type_utils.is_assignable_from(
expected_type, actual_type):
raise TypeError(
'Expected element at position {} to be '
'of type {}, found {}.'.format(
idx, str(expected_type), str(actual_type)))
if type_utils.is_anon_tuple_with_py_container(
element_value, expected_type):
element_value = type_utils.convert_to_py_container(
element_value, expected_type)
args.append(element_value)
kwargs = {}
for name, expected_type in six.iteritems(kwarg_types):
element_value = getattr(arg, name)
actual_type = type_utils.infer_type(element_value)
if not type_utils.is_assignable_from(
expected_type, actual_type):
raise TypeError('Expected element named {} to be '
'of type {}, found {}.'.format(
name, str(expected_type),
str(actual_type)))
if type_utils.is_anon_tuple_with_py_container(
element_value, expected_type):
element_value = type_utils.convert_to_py_container(
element_value, expected_type)
kwargs[name] = element_value
return fn(*args, **kwargs)
def _convert_to_py_container(value, type_spec):
"""Converts value to a Python container if type_spec has an annotation."""
if type_utils.is_anon_tuple_with_py_container(value, type_spec):
return type_utils.convert_to_py_container(value, type_spec)
elif isinstance(type_spec, computation_types.SequenceType):
if all(
type_utils.is_anon_tuple_with_py_container(
element, type_spec.element) for element in value):
return [
type_utils.convert_to_py_container(element, type_spec.element)
for element in value
]
return value
async def _invoke(executor, comp, arg):
"""A coroutine that handles invocation.
Args:
executor: An instance of `executor_base.Executor`.
comp: The first argument to `context_base.Context.invoke()`.
arg: The optional second argument to `context_base.Context.invoke()`.
Returns:
The result of the invocation.
"""
result_type = comp.type_signature.result
elements = [executor.create_value(comp)]
if isinstance(arg, anonymous_tuple.AnonymousTuple):
elements.append(executor.create_tuple(arg))
elements = await asyncio.gather(*elements)
comp = elements[0]
if len(elements) > 1:
arg = elements[1]
result = await executor.create_call(comp, arg)
result_val = _unwrap(await result.compute())
if type_utils.is_anon_tuple_with_py_container(result_val, result_type):
return type_utils.convert_to_py_container(result_val, result_type)
else:
return result_val
def test_is_anon_tuple_with_py_container(self):
self.assertTrue(
type_utils.is_anon_tuple_with_py_container(
anonymous_tuple.AnonymousTuple([('a', 0.0)]),
computation_types.NamedTupleTypeWithPyContainerType(
[('a', tf.float32)], dict)))
self.assertFalse(
type_utils.is_anon_tuple_with_py_container(
value_impl.ValueImpl(
computation_building_blocks.Data('nothing', tf.int32),
context_stack_impl.context_stack),
computation_types.NamedTupleTypeWithPyContainerType(
[('a', tf.float32)], dict)))
self.assertFalse(
type_utils.is_anon_tuple_with_py_container(
anonymous_tuple.AnonymousTuple([('a', 0.0)]),
computation_types.NamedTupleType([('a', tf.float32)])))
def _call(fn, parameter_type, arg):
arg_type = type_utils.infer_type(arg)
if not type_utils.is_assignable_from(parameter_type, arg_type):
raise TypeError('Expected an argument of type {}, found {}.'.format(
parameter_type, arg_type))
if type_utils.is_anon_tuple_with_py_container(arg, parameter_type):
arg = type_utils.convert_to_py_container(arg, parameter_type)
return fn(arg)
def invoke(self, fn, arg):
comp = self._compile(fn)
cardinalities = {}
root_context = ComputationContext(cardinalities=cardinalities)
computed_comp = self._compute(comp, root_context)
type_utils.check_assignable_from(comp.type_signature,
computed_comp.type_signature)
if not isinstance(computed_comp.type_signature,
computation_types.FunctionType):
if arg is not None:
raise TypeError('Unexpected argument {}.'.format(arg))
else:
value = computed_comp.value
result_type = fn.type_signature.result
if type_utils.is_anon_tuple_with_py_container(value, result_type):
return type_utils.convert_to_py_container(value, result_type)
return value
else:
if arg is not None:
def _handle_callable(fn, fn_type):
py_typecheck.check_type(fn, computation_base.Computation)
type_utils.check_assignable_from(fn.type_signature, fn_type)
computed_fn = self._compute(self._compile(fn), root_context)
return computed_fn.value
computed_arg = ComputedValue(
to_representation_for_type(arg,
computed_comp.type_signature.parameter,
_handle_callable),
computed_comp.type_signature.parameter)
cardinalities.update(
runtime_utils.infer_cardinalities(computed_arg.value,
computed_arg.type_signature))
else:
computed_arg = None
result = computed_comp.value(computed_arg)
py_typecheck.check_type(result, ComputedValue)
type_utils.check_assignable_from(comp.type_signature.result,
result.type_signature)
value = result.value
fn_result_type = fn.type_signature.result
if type_utils.is_anon_tuple_with_py_container(value, fn_result_type):
return type_utils.convert_to_py_container(value, fn_result_type)
return value
async def _invoke(executor, comp, arg):
"""A coroutine that handles invocation.
Args:
executor: An instance of `executor_base.Executor`.
comp: The first argument to `context_base.Context.invoke()`.
arg: The optional second argument to `context_base.Context.invoke()`.
Returns:
The result of the invocation.
"""
py_typecheck.check_type(comp.type_signature, computation_types.FunctionType)
result_type = comp.type_signature.result
if arg is not None:
py_typecheck.check_type(arg, executor_value_base.ExecutorValue)
comp = await executor.create_value(comp)
result = await executor.create_call(comp, arg)
py_typecheck.check_type(result, executor_value_base.ExecutorValue)
result_val = _unwrap(await result.compute())
if type_utils.is_anon_tuple_with_py_container(result_val, result_type):
return type_utils.convert_to_py_container(result_val, result_type)
else:
return result_val
def test_returns_false_with_named_tuple_type_spec(self):
value = anonymous_tuple.AnonymousTuple([('a', 0.0)])
type_spec = computation_types.NamedTupleType([('a', tf.float32)])
self.assertFalse(
type_utils.is_anon_tuple_with_py_container(value, type_spec))
def test_returns_false_with_none_value(self):
value = None
type_spec = computation_types.NamedTupleTypeWithPyContainerType(
[('a', tf.float32)], dict)
self.assertFalse(
type_utils.is_anon_tuple_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_utils.is_anon_tuple_with_py_container(value, type_spec))
