Find the first line in a multi-line import from statement with Python's AST

Question

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.

Asked By: DrownedSuccess

||

Source

Answer 1

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.)

Answered By: rici

Find the first line in a multi-line import from statement with Python's AST

Question:

Answers: