How to create a typed return using Python's ast? - `def f() -> bool: return True`

Viewed 215

How do I use the PEP484 syntax for return types, using the builtin ast library?

Attempt:

from ast import parse, Name, Load, Constant, Expr, FunctionDef, arguments, arg

func_def_ast = FunctionDef(
    args=arguments(
        args=[
            arg(annotation=Name(ctx=Load(), id='str'),
                arg='is_chocolate', type_comment=None),
            arg(annotation=Name(ctx=Load(), id='bool'),
                arg='verbose', type_comment=None)
        ],
        defaults=[Constant(kind=None, value=True)],
        kwarg=arg(annotation=None, arg='extra_kwargs',
                  type_comment=None), kw_defaults=[],
        kwonlyargs=[], posonlyargs=[], vararg=None
    ),
    body=[Expr(value=Constant(kind=None,
                              value='Crazy hungry function'))
    ],
    decorator_list=[], name='is_choco',
    returns=parse('Union[bool, str]'),
    type_comment=None
)

Now let's print it out using astor:

from astor import to_source    

src = to_source(func_def_ast)
print(src)

Which gives the following, on Python 3.8.3:

def is_choco(is_chocolate: str, verbose: bool=True, **extra_kwargs) ->
Union[bool, str]:
    """Crazy hungry function"""

Valid Python would require Union[bool, str] be on the same line that at the top, otherwise you get:

def is_choco(is_chocolate: str, verbose: bool=True, **extra_kwargs) ->
                                                                     ^

SyntaxError: invalid syntax

1 Answers

Ohhh I forgot to access it

This:

returns=parse('Union[bool, str]'),

Should be:

returns=parse('Union[bool, str]').body[0].value,
Related