I’m currently trying to find all ast.Import
and ast.ImportFrom
nodes in a Python file. However, in a multi-line import statement, like
from foo import ( bar, baz, coo )
the lineno mentioned for each submodule (bar
, baz
, and coo
) are all at the line where they are mentioned, not the original from foo import
line. How can I get just the line where the import statement begins (for single line imports, it would just be the same).
Alternatively, is there a way to get all imports in a scope, iteratively (going through all the scopes in the script)?
Update: Apparently, ast.walk
doesn’t return a node for each line, like I thought. It was actually because I made a different tuple for every node.names
. Fixing it to only return the first name (and use node.lineno) works out great. However @rici’s answer is still good, so I’ll leave this up.
Advertisement
Answer
You could just use the lineno
attribute of the ast.ImportFrom
node, which is the line number of the start of the statement. (There’s also an end_lineno
attribute, probably less useful for this case.)
Here’s a small example:
import ast sample = ''' # This is line 1 # # This is line 3 # # The next line is line 6 from foo import ( bar, baz, coo ) import pkg.mod, other_mod '''.strip() class myWalker(ast.NodeVisitor): def visit_Import(self, node): print(f"""Line {node.lineno} imports modules { ', '.join(alias.name for alias in node.names) }""") def visit_ImportFrom(self, node): print(f"""Line {node.lineno} imports from module {node.module} the names { ', '.join(alias.name for alias in node.names) }""") myWalker().visit(ast.parse(sample, filename='<test>', mode='exec'))
Output:
Line 6 imports from module foo the names bar, baz, coo Line 11 imports modules pkg.mod, other_mod
To get the import names by scope, I think you’d have to walk the syntax tree, noting function definitions. (Other scopes, such as lambdas and comprehensions, don’t matter here because they only allow expressions, not statements. But if you want to do it right, you’ll need to also keep track of names declared global
or nonlocal
, because a function could use one of those declarations to inject an imported name into a different scope.)