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