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.
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
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.
The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.