Assert that called computations are of functional type.
Now that no-argument lambdas are properly supported, functions
should all have a functional type, and this should be essentially
a no-op.
See cl/293639598 for more.
@@ -73,17 +73,20 @@ class FederatedComputationContext(context_base.Context):
def invoke(self, comp, arg):
fn = value_impl.to_value(comp, None, self._context_stack)
- if isinstance(fn.type_signature, computation_types.FunctionType):
+ tys = fn.type_signature
+ py_typecheck.check_type(tys, computation_types.FunctionType)
if arg is not None:
- type_utils.check_type(arg, fn.type_signature.parameter)
+ if tys.parameter is None:
+ raise ValueError(
+ 'A computation of type {} does not expect any arguments, but got '
+ 'an argument {}.'.format(tys, arg))
+ type_utils.check_type(arg, tys.parameter)
ret_val = fn(arg)
else:
+ if tys.parameter is not None:
+ raise ValueError(
+ 'A computation of type {} expects an argument of type {}, but got '
+ ' no argument.'.format(tys, tys.parameter))
ret_val = fn()
- type_utils.check_type(ret_val, fn.type_signature.result)
+ type_utils.check_type(ret_val, tys.result)
return ret_val
- elif arg is not None:
- raise ValueError(
- 'A computation of type {} does not expect any arguments, but got an '
- 'argument {}.'.format(fn.type_signature, arg))
- else:
- return fn
