给定对 function 的引用，在签名中生成不带注释的源代码

Question

Default values should remain.

Decorators should also be removed, but it's less of a concern.

Original source:

# Comments
@decorator1
@decorator2(
    a=1,
    b=1,
)
def my_func(
    a: typing.List = 14,
    b: 'CustomType' = None,
    c: Whatever('foo') = 42,
    d: Whatever('foo') = some_function,
) -> typing.NamedTuple(
    'Dummy',
    [
      ('r1': 'CustomType'),
      ('21': 'CustomType2'),
    ]
):
    ...

Desired source:

def my_func(
    a = 14,
    b = None,
    c = 42,
    d = some_function,
):
    ...

I can get the source using inspect.getsource(my_func) , but I need to strip the annotations. How can I do this? Preferably using the standard modules.

Answer 1

You can use ast.parse to parse the source into an AST, and use ast.walk to traverse the tree to nullify annotation if it's an arg node and nullify returns and decorator_list if it's a FunctionDef node. Use astunparse to unparse the tree back to source code:

import inspect
import ast
import typing
import astunparse
from unittest.mock import patch

@patch('builtins.print')
def my_func(
        a: int = 1,
        b: typing.List = []
) -> bool:
    pass

tree = ast.parse(inspect.getsource(my_func), '', 'exec')
for node in ast.walk(tree):
    if isinstance(node, ast.arg):
        node.annotation = None
    elif isinstance(node, ast.FunctionDef):
        node.returns = None
        node.decorator_list = []
print(astunparse.unparse(tree))

This outputs:

def my_func(a=1, b=[]):
    pass

Demo: https://repl.it/repls/WaterloggedFunnyQuotient

Answer 2

You can subclass lib2to3.refactor.RefactoringTool to refactor the code using a fixer that is a subclass of lib2to3.fixer_base.BaseFix with a pattern that looks for either a typed argument, a function declaration with an annotated returning value, or a decorated definition, and a transform method that removes the indices of the annotations and decorators from the child nodes:

from lib2to3 import fixer_base, refactor

class FixParameterAnnotations(fixer_base.BaseFix):
    PATTERN = "name=tname | func=funcdef< any+ '->' any+ > | decorated"

    def transform(self, node, results):
        if 'name' in results:
            del node.children[1:] # delete annotation to typed argument
        elif 'func' in results:
            del node.children[-4:-2] # delete annotation to function declaration
        else:
            del node.children[0] # delete decorators
        return node

class Refactor(refactor.RefactoringTool):
    def __init__(self, fixers):
        self._fixers= [cls(None, None) for cls in fixers]
        super().__init__(None)

    def get_fixers(self):
        return self._fixers, []

so that:

import inspect
import typing
from unittest.mock import patch

@patch('builtins.print')
def my_func(
        a: int = 1, # some comment
        b: typing.List = []     # more comment
) -> bool:
    ''' some docstring'''
    pass

print(Refactor([FixParameterAnnotations]).refactor_string(inspect.getsource(my_func), ''))

outputs:

def my_func(
        a = 1, # some comment
        b = []     # more comment
):
    ''' some docstring'''
    pass

Demo: https://repl.it/@blhsing/BrightWhirlwindBoolean

lib2to3 is round-trip stable so all comments and white spaces are preserved after the transformation. You can find the definition of the Python grammar in Grammar.txt of the lib2to3 module.