Aluode/PerceptionLabPortable

1"""
2This module provides helper functions to find the first line of a function
3body.
4"""
5
6import ast
7import inspect
8import textwrap
9
10
11class FindDefFirstLine(ast.NodeVisitor):
12    """
13    Attributes
14    ----------
15    first_stmt_line : int or None
16        This stores the first statement line number if the definition is found.
17        Or, ``None`` if the definition is not found.
18    """
19
20    def __init__(self, name, firstlineno):
21        """
22        Parameters
23        ----------
24        code :
25            The function's code object.
26        """
27        self._co_name = name
28        self._co_firstlineno = firstlineno
29        self.first_stmt_line = None
30
31    def _visit_children(self, node):
32        for child in ast.iter_child_nodes(node):
33            super().visit(child)
34
35    def visit_FunctionDef(self, node: ast.FunctionDef):
36        if node.name == self._co_name:
37            # Name of function matches.
38
39            # The `def` line may match co_firstlineno.
40            possible_start_lines = set([node.lineno])
41            if node.decorator_list:
42                # Has decorators.
43                # The first decorator line may match co_firstlineno.
44                first_decor = node.decorator_list[0]
45                possible_start_lines.add(first_decor.lineno)
46            # Does the first lineno match?
47            if self._co_firstlineno in possible_start_lines:
48                # Yes, we found the function.
49                # So, use the first statement line as the first line.
50                if node.body:
51                    first_stmt = node.body[0]
52                    if _is_docstring(first_stmt):
53                        # Skip docstring
54                        first_stmt = node.body[1]
55                    self.first_stmt_line = first_stmt.lineno
56                    return
57                else:
58                    # This is probably unreachable.
59                    # Function body cannot be bare. It must at least have
60                    # A const string for docstring or a `pass`.
61                    pass
62        self._visit_children(node)
63
64
65def _is_docstring(node):
66    if isinstance(node, ast.Expr):
67        if (isinstance(node.value, ast.Constant)
68                and isinstance(node.value.value, str)):
69            return True
70    return False
71
72
73def get_func_body_first_lineno(pyfunc):
74    """
75    Look up the first line of function body using the file in
76    ``pyfunc.__code__.co_filename``.
77
78    Returns
79    -------
80    lineno : int; or None
81        The first line number of the function body; or ``None`` if the first
82        line cannot be determined.
83    """
84    co = pyfunc.__code__
85    try:
86        with open(co.co_filename) as fin:
87            source = fin.read()
88            offset = 0
89    except (FileNotFoundError, OSError):
90        try:
91            lines, offset = inspect.getsourcelines(pyfunc)
92            source = "".join(lines)
93            offset = offset - 1
94        except (OSError, TypeError):
95            return None
96
97    tree = ast.parse(textwrap.dedent(source))
98    finder = FindDefFirstLine(co.co_name, co.co_firstlineno - offset)
99    finder.visit(tree)
100    if finder.first_stmt_line:
101        return finder.first_stmt_line + offset
102    else:
103        # No first line found.
104        return None
105