Search code examples
python-3.xcompilationabstract-syntax-treepython-ast

Python: what is the ast.FunctionType in python ast?


I don't pretend to have understood everything in python ast, but FunctionType is something that bothered me.

mod = Module(stmt* body, type_ignore *type_ignores)
        | Interactive(stmt* body)
        | Expression(expr body)
        | FunctionType(expr* argtypes, expr returns)
        | Suite(stmt* body) 

What can it be?


Solution

  • FunctionType is the top-level symbol returned for parsing # type: (args) -> return_type comments for function types (as specified by PEP484)

    The only authoritative source for this I can find is the changelog for Python 3.8 where it was added: https://docs.python.org/3/whatsnew/3.8.html#ast

    You can parse this type of thing with compile or ast.parse with the 'func_type' mode (where the mode would usually be exec for statements and eval for expressions).

    This is only used (by external type checking software) if you are trying to annotate function return types with the Python 2 syntax, like:

    # Invalid syntax in Python 2
    # def f(x: int, y: int) -> int:
    #     return x * y
    
    # Special type of type hint for functions in Python 2  (and also works in Python 3)
    def f(x, y):
        # type: (int, int) -> int
        return x * y
    
    # Or alternative syntax:
    def f(
        x,  # type: int
        y   # type: int
    ):
        # type: (...) -> int
        return x * y
    

    Normal type annotations can just be parsed in 'eval' mode, but function type annotations can't be, so a new mode and top-level symbol was added:

    https://github.com/python/cpython/blob/main/Grammar/python.gram#L91

    func_type[mod_ty]: '(' a=[type_expressions] ')' '->' b=expression NEWLINE* ENDMARKER { _PyAST_FunctionType(a, b, p->arena) }
    

    Which can be used by type checkers to parse these types of comments:

    import ast
    
    source = '''
    def f(x, y):
        # type: (int, int) -> int
        return x * y
    '''
    
    module = ast.parse(source, mode='exec', type_comments=True)
    
    class PrintFunctionTypes(ast.NodeVisitor):
        def visit_FunctionDef(self, f):
            print('FunctionDef:\n' + ast.dump(f))
            print('\ntype_comment:\n' + f.type_comment)
            parsed_comment = ast.parse(f.type_comment, mode='func_type')
            print('\nWhen parsed:\n' + ast.dump(parsed_comment))
    
    PrintFunctionTypes().visit(module)
    

    Which outputs:

    FunctionDef:
    FunctionDef(name='f', args=arguments(posonlyargs=[], args=[arg(arg='x'), arg(arg='y')], kwonlyargs=[], kw_defaults=[], defaults=[]), body=[Return(value=BinOp(left=Name(id='x', ctx=Load()), op=Mult(), right=Name(id='y', ctx=Load())))], decorator_list=[], type_comment='(int, int) -> int')
    
    type_comment:
    (int, int) -> int
    
    When parsed:
    FunctionType(argtypes=[Name(id='int', ctx=Load()), Name(id='int', ctx=Load())], returns=Name(id='int', ctx=Load()))