如何获取Python导入模块的行号与列偏移信息？

我有一个基于ast.NodeVisitor的类，用于提取Python文件中的导入模块列表，现在想同时获取模块名称在文件中的行号和列偏移量，现有代码、当前输出及期望输出如下，请问该如何修改代码？

现有代码 #

import ast

class ImportFinder(ast.NodeVisitor):
    def __init__(self):
        self.imports = []

    def visit_Import(self, node):
        for i in node.names:
            self.imports.append({'import_type': "import", 'module': i.name,})

    def visit_ImportFrom(self, node):
        self.imports.append({'import_type': "from", 'module': node.module})

def parse_imports(source):
    tree = ast.parse(source)
    finder = ImportFinder()
    finder.visit(tree)
    return finder.imports


# Example usage
sample_file = '''
from foo import bar, baz, frob
import bar.baz
import   bar.foo as baf
'''
parsed_imports = parse_imports(sample_file)
for i in parsed_imports:
    print(i)

当前输出 #

{'import_type': 'from', 'module': 'foo'}
{'import_type': 'import', 'module': 'bar.baz'}
{'import_type': 'import', 'module': 'bar.foo'}

期望输出 #

{'import_type': 'from', 'module': 'foo', 'line': 2, 'column_offset': 5}
{'import_type': 'import', 'module': 'bar.baz', 'line': 3, 'column_offset': 7}
{'import_type': 'import', 'module': 'bar.foo', 'line': 4, 'column_offset': 9}

解决方法 #

Python AST节点自带lineno（行号）和col_offset（列偏移）属性，我们可以利用这些属性直接获取模块名称的位置信息：

• 对于Import节点：每个导入的模块名对应ast.alias对象，直接取该对象的lineno和col_offset即可得到模块名称的行号与起始列偏移。
• 对于ImportFrom节点：模块名称的行号等于ImportFrom节点的lineno，列偏移则是ImportFrom节点的col_offset加上'from '的长度（即5），因为col_offset是from关键字的起始位置，加上from 的长度后正好是模块名称的起始位置。

修改后的完整代码如下：

import ast

class ImportFinder(ast.NodeVisitor):
    def __init__(self):
        self.imports = []

    def visit_Import(self, node):
        for i in node.names:
            self.imports.append({
                'import_type': "import",
                'module': i.name,
                'line': i.lineno,
                'column_offset': i.col_offset
            })

    def visit_ImportFrom(self, node):
        module_col_offset = node.col_offset + len('from ')
        self.imports.append({
            'import_type': "from",
            'module': node.module,
            'line': node.lineno,
            'column_offset': module_col_offset
        })

def parse_imports(source):
    tree = ast.parse(source)
    finder = ImportFinder()
    finder.visit(tree)
    return finder.imports


# Example usage
sample_file = '''
from foo import bar, baz, frob
import bar.baz
import   bar.foo as baf
'''
parsed_imports = parse_imports(sample_file)
for i in parsed_imports:
    print(i)

运行后即可得到符合期望的输出。

内容的提问来源于stack exchange，提问作者Stevoisiak