def update_signature_return(
sig: inspect.Signature,
return_type: type = None,
yield_type: type = None,
existing_annotation_strategy: ExistingAnnotationStrategy = ExistingAnnotationStrategy.REPLICATE,
) -> inspect.Signature:
"""Update return annotation with the supplied types"""
anno = sig.return_annotation
if anno is not inspect.Signature.empty:
# If generating a stub to apply and there's already a return type
# annotation, generate a stub with no return type annotation, to avoid
# the possibility of "incompatible annotation" errors.
if existing_annotation_strategy == ExistingAnnotationStrategy.OMIT:
return sig.replace(return_annotation=inspect.Signature.empty)
# Don't change pre-existing annotations unless asked to
if existing_annotation_strategy == ExistingAnnotationStrategy.REPLICATE:
return sig
# NB: We cannot distinguish between functions that explicitly only
# return None and those that do so implicitly. In the case of generator
# functions both are typed as Iterator[<yield_type>]
if (yield_type is not None) and ((return_type is None) or (return_type == NoneType)):
anno = make_iterator(yield_type)
elif (yield_type is not None) and (return_type is not None):
anno = make_generator(yield_type, NoneType, return_type)
elif return_type is not None:
anno = return_type
return sig.replace(return_annotation=anno)
def update_signature_return(sig: inspect.Signature,
return_type: type = None,
yield_type: type = None) -> inspect.Signature:
"""Update return annotation with the supplied types"""
anno = sig.return_annotation
# Dont' touch pre-existing annotations
if anno is not inspect.Signature.empty:
return sig
# NB: We cannot distinguish between functions that explicitly only
# return None and those that do so implicitly. In the case of generator
# functions both are typed as Iterator[<yield_type>]
if (yield_type is not None) and ((return_type is None) or
(return_type == NoneType)):
anno = make_iterator(yield_type)
elif (yield_type is not None) and (return_type is not None):
anno = make_generator(yield_type, NoneType, return_type)
elif return_type is not None:
anno = return_type
return sig.replace(return_annotation=anno)
