CODEWITHSUNDEEP
pythonabstract-syntax-tree
tiggerae
tiggerae

Reputation: 1501

How do I list a function's parameters using ast?

I'm using the ast module to parse docstrings in a Python module to turn our docs into read the docs format. I'm using the following to get the function names and docstrings into a list of dicts which the rest of my code works well with. I'm looking for something to get the parameters of the function as well:

good_file = (file for file in os.listdir() if file[-3:] == '.py' and file != '__init__.py')

functions = []

for file in good_file:
     with open(file, 'r') as f:
         module = ast.parse(f.read())

     for node in module.body:
         if isinstance(node, ast.FunctionDef):
         entry = {"docs": ast.get_docstring(node), "fn_name": node.name, "params": ???}

         functions.append(entry)

I'm looking for what I can use to fill in the function's parameters into the dict. Thanks!

Upvotes: 6

Views: 7818

Answers (1)

Martijn Pieters
Martijn Pieters

Reputation: 1125268

The Abstract Grammar section of the ast documentation tells you where to find the parameter definitions in a FunctionDef node:

stmt = FunctionDef(identifier name, arguments args,
                   stmt* body, expr* decorator_list, expr? returns)

In the parameters is a sequence of type name entries; the names become attributes on the node. The types are further covered in the documentation (with several 'builtin' types listed at the top, which are reflected as Python strings and integers). Types with * after them are sequences (lists), a question mark means they can be set to None.

So each FunctionDef node has name, args, body, decorator_list and returns attributes. The args attribute is a new node, of type arguments, also documented:

arguments = (arg* args, arg? vararg, arg* kwonlyargs, expr* kw_defaults,
             arg? kwarg, expr* defaults)

so FunctionDef.args.args is a list of arguments, each an arg object, etc.

arg is documented as

arg = (identifier arg, expr? annotation)
       attributes (int lineno, int col_offset)

where identifier is a built-in type, so just a string here.

You probably want to look at the ast.dump() function, which will give you a quick overview of an AST node:

>>> source = """def foo(bar, baz=None, *args, **kwargs): pass"""
>>> module = ast.parse(source)
>>> ast.dump(module)
"Module(body=[FunctionDef(name='foo', args=arguments(args=[arg(arg='bar', annotation=None), arg(arg='baz', annotation=None)], vararg=arg(arg='args', annotation=None), kwonlyargs=[], kw_defaults=[], kwarg=arg(arg='kwargs', annotation=None), defaults=[NameConstant(value=None)]), body=[Pass()], decorator_list=[], returns=None)])"

From there you can then explore and further 'dump' information to get to the actual data you need:

>>> function = module.body[0]
>>> ast.dump(function.args)
"arguments(args=[arg(arg='bar', annotation=None), arg(arg='baz', annotation=None)], vararg=arg(arg='args', annotation=None), kwonlyargs=[], kw_defaults=[], kwarg=arg(arg='kwargs', annotation=None), defaults=[NameConstant(value=None)])"
>>> function.args.args
[<_ast.arg object at 0x109852fd0>, <_ast.arg object at 0x109852ef0>]
>>> [a.arg for a in function.args.args]
['bar', 'baz']

The defaults are attached to the last names in the args or kw_args sequences (defaults for args, kw_defaults for kwonlyargs); a list of N defaults attached to the last N names in args or kwosnlyargs. Any catch-all names (*args and **kwargs in my example) are listed separately.

Upvotes: 15

Related Questions