In Python, given a line number, how do I know which scope it belongs to?

user_2142577 0
Viewed 354

Say I have a simple program:

1  def foo():
2     pass
3
4  def bar():
5     pass

Given line number 1 or 2, I want to know it belongs to function foo.

Same with 4 or 5 for function bar.

What's the easiest way to do it? For simplicity, let's just assume "scope" means function or method.

I'm looking at libCST as it provides the code position metadata, but it does not provide a mapping from line number to its scope.

python
2 Answers

You can use ast to walk over the file structure, determine the line ranges of each function and then find where your specified lineno resides.

import ast

def find_definition(linenum):
    functions = {}
    with open('input.py') as file:
        tree = ast.parse(file.read())

        for item in ast.walk(tree):
            if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
                start, end = compute_size(item)

                functions[item.name] = (start, end)

    for key, value in functions.items():
        if value[0] <= linenum <= value[1]:
            return key, value

def compute_size(node):
    min_lineno = node.lineno
    max_lineno = node.lineno
    for node in ast.walk(node):
        if hasattr(node, "lineno"):
            min_lineno = min(min_lineno, node.lineno)
            max_lineno = max(max_lineno, node.lineno)
    return (min_lineno, max_lineno + 1)

function, func_range = find_definition(1)

print(function, func_range)

#foo (1, 3)

function, func_range = find_definition(4)

print(function, func_range)

#bar (4, 6)

There is a much more indepth read found here.

user_79428560

The standard library ast module will let you parse the file the same way that (your current version of) Python sees it. That will give you a data structure with line number information (lineno and end_lineno attributes) for each node in your code, as well as columns if you need them.

Then you can go through the data structure in as much detail as you want, mapping line numbers to the corresponding functions.

user_6833290
Related