Files
renderdoc/qrenderdoc/Code/pyrenderdoc/parse_reflection.py
T
baldurk e57d9db9c4 Process scope declarations and identifiers in AST tree
* We don't generate type information yet, _get_type is stubbed out.
2026-08-13 17:46:41 +01:00

728 lines
27 KiB
Python

import ast
import builtins
from typing import List, Dict, Any, Tuple, Callable, TypeVar, Optional
# return whether an object is a specialisation of the given generic,
# e.g. _is_generic(Dict, Dict[str, int]) == True
def _is_generic(generic, obj):
if hasattr(obj, "__origin__") and hasattr(generic, "__origin__"):
if generic.__origin__ == obj.__origin__:
return True
if hasattr(obj, "__origin__"):
return obj.__origin__ == generic
return False
# return true for AST nodes that need their own scope - this is module level, then
# classes and functions which may be nested inside each other
def _is_scope_node(node):
return (
isinstance(node, ast.Module)
or isinstance(node, ast.ClassDef)
or isinstance(node, ast.FunctionDef)
or isinstance(node, ast.AsyncFunctionDef)
)
# get all the return statements immediately inside a function without going into nested
# functions
def _get_return_statements(node: ast.AST, root: bool = True):
if _is_scope_node(node) and not root:
return []
if isinstance(node, ast.Return):
return [node]
ret = []
for recurse in ["body", "orelse", "finalbody"]:
if hasattr(node, recurse):
for n in getattr(node, recurse):
if not _is_scope_node(n):
ret += _get_return_statements(n, False)
return ret
# Python 3.8+ is expected to have start and end lines.
# Before that, end was missing so we assume all statements are single
# line (or the last child).
def _get_linerange(node: ast.AST):
if not hasattr(node, "lineno"):
return (-1, -1)
lineno = getattr(node, "lineno")
if hasattr(node, "end_lineno"):
return (lineno, getattr(node, "end_lineno"))
end_lineno = lineno
for recurse in ["body", "orelse", "finalbody"]:
if hasattr(node, recurse) and len(getattr(node, recurse)) > 0:
end_lineno = max(end_lineno, _get_linerange(getattr(node, recurse)[-1])[1])
return (lineno, end_lineno)
# true if this is a `self.foo` type attribute lookup, so we know
# to look up in the parent scope which is how we handle self
def _is_self_lookup(parsed: ast.AST):
return (
isinstance(parsed, ast.Attribute)
and isinstance(parsed.value, ast.Name)
and parsed.value.id == "self"
)
# comment out all lines starting from a given point,
# to try and make things compile. Stops when it hits
# an indent that looks like the end of the statement
def _commentlines(text, first_comment_line):
lines = text.splitlines()
indent = len(lines[first_comment_line]) - len(lines[first_comment_line].lstrip())
lines[first_comment_line] = "#" + lines[first_comment_line]
for i in range(first_comment_line + 1, len(lines)):
if lines[i].strip() == "":
continue
if lines[i].startswith(" " * (indent + 1)):
lines[i] = "#" + lines[i]
continue
break
return "\n".join(lines)
# a given instance of an identifier and its type
class Ident:
# first line this ident is valid
line: int = -1
# for assignments the identifier is only valid on the first line
# before a certain column.
# mostly relevant for overwriting assignments e.g. foo = foo.bar
# so the idents for LHS foo and RHS foo can be differentiated
col: int = 9999
# the type or type hint
type_obj: Optional[Any] = None
# for functions without return annotations, this will be set to the AST node
# for lazy evaluation to obtain a guessed return type.
# We do it this way as we normally process in declaration order but
lazy_node: Optional[ast.AST] = None
# a scope - either a module, class or function
class Scope:
# the name for debugging
name: str
# the parent scope, for searches upwards for identifiers
parent: "Optional[Scope]" = None
# the parsed node
parsed: ast.AST
# the type, only relevant for classes
type_obj: Optional[Any] = None
# known identifiers in this scope
identifiers: Dict[str, List[Ident]]
# for non-modules, the ident of this scope
ident: Optional[Ident] = None
# whether this scope is a class or not (for finding `self`)
is_class: bool = False
def __init__(self):
self.identifiers = {}
def set_ident(self, name: str, ident: Ident):
if name not in self.identifiers:
self.identifiers[name] = []
self.identifiers[name] += [ident]
# look up the version of an identifier on a given line, in our parents,
# or in the builtins
def get_ident(self, name: str, line: int, col: int):
ret = None
if name in self.identifiers:
for i in self.identifiers[name]:
# only consider identifiers that are valid for the line & col
# we're searching for
if (
i.line < line
or (i.line == line and (col < i.col or col == -1))
or line == -1
):
# if we don't have a match, or this match is more recent, use it
if ret is None or ret.line < i.line:
ret = i
# if we don't have a record of it, or we're at a statement before
# the first assignment, search the parent at our declaration line
if ret is None:
if self.parent is not None:
return self.parent.get_ident(
name, self.ident.line if self.ident is not None else line, col
)
if name in dir(builtins):
return getattr(builtins, name)
return None
return ret
def full_name(self):
if self.parent is not None:
return f"{self.parent.full_name()}::{self.name}"
return self.name
def __repr__(self):
return f"<Scope '{self.full_name()}'>"
# Main class, reflects a given source text (if it can) and allows
# lookups of the types of expressions as well as auto-completion
# of partial expressions
class PyReflector:
def __init__(self, text: str, starting_globals: Dict[str, Any], debug_types: bool):
# the parsed module, or None if parsing completely failed
self.module: Optional[ast.Module]
# for tests - a lookup to retrieve the actual typevar since they can't be compared by name
self.user_types: Dict[str, TypeVar] = {}
# the text that was actually parsed, including any truncation/commenting needed
# to get it to compile
self.parsed_text: str
# an error if parsing completely failed
self.parse_error: Optional[str]
# the starting set of globals to consider the module populated with
self.starting_globals = starting_globals
if self.starting_globals is None:
self.starting_globals = globals()
# whether or not type-processing should be debugged. Instead of falling back
# to `typing.Any` for unknown types, instead a string bounded by "@@" is returned.
# Mostly for internal use
self.debug_types = debug_types
# the current scope for each line
self.scopes: List[Scope] = []
# a pending FIFO which we process classes and functions in. This is used
# so that if we want to infer a type we can steal it early and process functions
# out of normal declaration order. Can't resolve mutually-recursive functions
# that require type guessing but improves many common situations where a class
# method calls another that's declared later and doesn't have proper type
# annotations
self.pending: List[Tuple[Optional[Scope], ast.AST]] = []
# try to parse the text
self._parse_text(text)
# This checks module against None internally to help type checkers
self._process_module()
def _parse_text(self, text: str):
# try a simple parse. If there are no syntax errors this will
# succeed.
try:
self.module = ast.parse(text)
self.parsed_text = text
self.parse_error = None
return
except SyntaxError as err:
if err.lineno is None:
raise err
first_comment_line = err.lineno - 1
# when encountering an error, comment everything
# from the error line to the next line with same or
# less indent (excluding blank lines) and try again
mod = _commentlines(text, first_comment_line)
try:
self.module = ast.parse(mod)
self.parsed_text = mod
self.parse_error = None
return
except SyntaxError as err2:
if err2.lineno is None:
raise err2
# if the error has moved to a later line, that suggests the
# original error was reported from some previous line, so we
# should try from an earlier point
# if not, we can't recover this
if err2.lineno <= first_comment_line + 1:
self.module = None
self.parsed_text = mod
self.parse_error = "Error remained after comments"
# first see where we can truncate to and successfully parse
# (up to 10 lines of non-blank lines truncated, to limit scope)
lines = text.splitlines()
trunc_lines = lines[0 : first_comment_line + 1]
removed = 0
while len(trunc_lines) > 0:
if trunc_lines[-1].strip() == "":
del trunc_lines[-1]
continue
del trunc_lines[-1]
removed += 1
if len(trunc_lines) == 0:
break
if trunc_lines[-1].rstrip()[-1] == ":":
trunc_lines[-1] += " pass"
lines[len(trunc_lines) - 1] += " pass"
try:
parsed = ast.parse("\n".join(trunc_lines))
text = "\n".join(lines)
break
except Exception:
if removed >= 10:
self.module = None
self.parsed_text = "\n".join(trunc_lines)
self.parse_error = "Couldn't backtrack"
# now we know that trunc_lines parses,
# retry commenting starting from there
mod = _commentlines(text, len(trunc_lines))
try:
self.module = ast.parse(mod)
self.parsed_text = mod
self.parse_error = None
return
except Exception:
self.module = None
self.parsed_text = mod
self.parse_error = "Error remained after comments"
def valid(self):
return self.module is not None
# get the source string for a given line
def get_line_source(self, line: int):
return self.parsed_text.splitlines()[line - 1]
def _type_failure(self, err: str):
if self.debug_types:
return "@@" + err.replace("@", "") + "@@"
return Any
def _get_type(self, scope: Scope, parsed: Optional[ast.AST]):
return None
# we can try to guess function return values by looking at the types of
# the return statements. If they are all the same (ignoring possible none)
# then use that. If they're different we give up as we don't handle union/
# varied types
def _guess_return_value(self, scope: Scope, node: ast.AST):
ret_types = [
self._get_type(scope, r.value) for r in _get_return_statements(node)
]
ret_types = list(set([r for r in ret_types if r is not None]))
if len(ret_types) == 1 and not isinstance(ret_types[0], str):
return ret_types[0]
return type(None)
# for a single statement, process the identifiers it creates and recurse
# as needed (but not into new scopes)
def _process_stmt(self, parent: Optional[Scope], parsed: ast.AST):
if parent is None:
raise ValueError("Expected parent for non-module")
# for functions and classes we register their type as an identifier and update
# scopes, but push them onto the pending list and continue processing.
if isinstance(parsed, ast.ClassDef):
classscope = Scope()
classscope.name = f"class {parsed.name}"
classscope.parent = parent
classscope.parsed = parsed
classscope.type_obj = TypeVar(parsed.name) # type: ignore
self.user_types[parsed.name] = classscope.type_obj
classscope.is_class = True
r = _get_linerange(parsed)
for line in range(r[0], r[1] + 1):
self.scopes[line] = classscope
id = Ident()
id.line = parsed.lineno
id.type_obj = classscope.type_obj
classscope.ident = id
parent.set_ident(parsed.name, id)
self.pending.append((classscope, parsed))
return
elif isinstance(parsed, ast.FunctionDef) or isinstance(
parsed, ast.AsyncFunctionDef
):
funcscope = Scope()
funcscope.name = f"function {parsed.name}"
funcscope.parent = parent
funcscope.parsed = parsed
funcscope.is_class = False
r = _get_linerange(parsed)
for line in range(r[0], r[1] + 1):
self.scopes[line] = funcscope
id = Ident()
id.line = parsed.lineno
id.type_obj = self._get_type(parent, parsed)
# if there's no return annotation, we'll try to guess it
# later when this function gets processed
if parsed.returns is None:
id.lazy_node = parsed
else:
# if we know the return type already and this is a @property
# then pretend it is just a member of that type, not a function
if any(
[
isinstance(a, ast.Name) and a.id == "property"
for a in parsed.decorator_list
]
):
id.type_obj = self._get_type(parent, parsed.returns)
prop_setter = False
# if this one is a property setter, @self.setter, then
# don't add it as an ident
if any(
[
isinstance(a, ast.Attribute)
and a.attr == "setter"
and isinstance(a.value, ast.Name)
and a.value.id == parsed.name
for a in parsed.decorator_list
]
):
prop_setter = True
funcscope.ident = id
if not prop_setter:
parent.set_ident(parsed.name, id)
args = parsed.args
arg_list = []
if "posonlyargs" in args._fields:
arg_list += args.posonlyargs
arg_list += args.args
# resize up the defaults array to the right size, defaults are 'trailing'
# ie. if there are fewer defaults than arguments, the first ones (starting
# from position-only arguments) have defaults omitted
defaults = [None] * (len(arg_list) - len(args.defaults)) + args.defaults
# len(kwonlyargs) == len(kw_defaults) because keyword defaults can come in
# any order
arg_list += args.kwonlyargs
defaults += args.kw_defaults
for i, a in enumerate(arg_list):
id = Ident()
id.line = parsed.lineno
default_val = defaults[i]
if a.annotation is not None:
id.type_obj = self._get_type(parent, a.annotation)
elif default_val is not None:
id.type_obj = self._get_type(parent, default_val)
elif (
a in parsed.args.args
and parsed.args.args.index(a) == 0
and a.arg == "self"
and parent.parent is not None
):
id.type_obj = parent.type_obj
else:
if self.debug_types:
id.type_obj = self._type_failure(
f"Unknown parameter type {a.arg} in {parsed.name}"
)
id.type_obj = Any
funcscope.set_ident(a.arg, id)
self.pending.append((funcscope, parsed))
return
# for imports, try to import the module ourselves so that we can have proper types.
# this won't work well for relative imports or things that need a particular sys.path
# but will work for standard library modules
if isinstance(parsed, ast.Import):
for alias in parsed.names:
n = alias.asname
if n is None or n == "":
n = alias.name
id = Ident()
id.line = parsed.lineno
try:
id.type_obj = __import__(alias.name, globals(), locals())
except ImportError:
if self.debug_types:
print(f"Couldn't import {alias.name}")
id.type_obj = Any
parent.set_ident(n, id)
if isinstance(parsed, ast.ImportFrom):
module = None
if parsed.module is not None:
try:
module = __import__(
parsed.module,
globals(),
locals(),
[a.name for a in parsed.names],
parsed.level,
)
except ImportError:
if self.debug_types:
print(f"Couldn't import {parsed.module}")
module = None
for alias in parsed.names:
n = alias.asname
if n is None or n == "":
n = alias.name
id = Ident()
id.line = parsed.lineno
if module is None:
try:
id.type_obj = __import__(
alias.name, globals(), locals(), [], parsed.level
)
except ImportError:
if self.debug_types:
print(f"Couldn't import {alias.name}")
id.type_obj = Any
else:
if hasattr(module, alias.name):
id.type_obj = getattr(module, alias.name)
else:
id.type_obj = Any
parent.set_ident(n, id)
# global/nonlocal we ignore for now, we assume the type won't
# change with any assignments there
targets = []
values = []
unpacking = False
ident_col = None
# AugAssign doesn't create a new object, so ignore it
# for other things that create a new identifier register both the set of
# target names and the source values. For pure assignments note the
# column where the LHS ends so that we can identify both possibly different
# types of `foo` in the statement `foo = foo.bar`
if isinstance(parsed, ast.For) or isinstance(parsed, ast.AsyncFor):
if isinstance(parsed.target, ast.Tuple) or isinstance(
parsed.target, ast.List
):
targets = parsed.target.elts
else:
targets = [parsed.target]
values = [parsed.iter] * len(targets)
unpacking = True
# unless this actually creates a new name we don't have to do anything
if isinstance(parsed, ast.With) or isinstance(parsed, ast.AsyncWith):
for item in parsed.items:
if item.optional_vars is not None:
if isinstance(item.optional_vars, ast.Name):
targets += [item.optional_vars]
values += [item.context_expr]
elif isinstance(item.optional_vars, ast.Tuple):
unpacking = True
targets += item.optional_vars.elts
values += [item.context_expr] * len(item.optional_vars.elts)
# for annotated assignments, trust the annotation and don't try to evaluate the
# actual RHS
if isinstance(parsed, ast.AnnAssign):
targets += [parsed.target]
values += [parsed.annotation]
if parsed.value is not None:
ident_col = parsed.value.col_offset
if isinstance(parsed, ast.Assign):
if len(parsed.targets) == 1 and (
isinstance(parsed.targets[0], ast.Tuple)
or isinstance(parsed.targets[0], ast.List)
):
targets = parsed.targets[0].elts
unpacking = True
else:
targets = parsed.targets
values = [parsed.value] * len(targets)
if parsed.value is not None:
ident_col = parsed.value.col_offset
# starred has no effect for our purposes
for i in range(len(targets)):
t = targets[i]
while isinstance(t, ast.Starred):
t = t.value
targets[i] = t
# we were in control of these arrays so they should be identically sized
if len(targets) != len(values):
raise RuntimeError("Didn't get equal number of targets and values")
for i in range(len(targets)):
t = targets[i]
ident_scope = parent
if isinstance(t, ast.Name):
name = t.id
if isinstance(t.ctx, ast.Load):
raise ValueError("Didn't expect loading target")
elif _is_self_lookup(t) and parent.parent is not None:
# just to help the type checked, is_self_lookup already checked this
if isinstance(t, ast.Attribute):
name = t.attr
ident_scope = parent.parent
# walk up to the class, in case of nested functions
while (
ident_scope.type_obj is None and ident_scope.parent is not None
):
ident_scope = ident_scope.parent
if isinstance(t.ctx, ast.Load):
raise ValueError("Didn't expect loading target")
else:
raise RuntimeError("invalid")
else:
# otherwise do nothing, this is an assignment of a value and we don't
# track fully dynamic types and attributes
continue
id = Ident()
id.line = _get_linerange(parsed)[0]
v = values[i]
id.type_obj = self._get_type(parent, v)
if unpacking:
t: Any = id.type_obj
if _is_generic(Tuple, t):
# either out of bounds, or a `Tuple[...]`, either way call it Any
if i < len(t.__args__):
id.type_obj = t.__args__[i]
else:
id.type_obj = Any
elif _is_generic(List, t):
id.type_obj = t.__args__[0]
else:
id.type_obj = t
if ident_col is not None:
id.col = ident_col
ident_scope.set_ident(name, id)
# should only get here for things like loops, ifs, etc NOT for classes and functions
if _is_scope_node(parsed):
raise TypeError("Should not be recursing for scope node")
for recurse in ["body", "orelse", "finalbody"]:
if recurse in parsed._fields:
for e in getattr(parsed, recurse):
self._process_stmt(parent, e)
# function to process the n'th item in the pending list. Usually 0 to
# continue processing in declaration order but can be out-of-order if we
# want to crystallise a guessed return type for a function in order to
# get a better type at an earlier callsite
def _process_pending(self, idx: int):
scope, node = self.pending.pop(idx)
# node is a module, class, or function. Process all the identifiers in it
# and add any nested classes or functions to the pending list
if _is_scope_node(node):
for st in getattr(node, "body"):
self._process_stmt(scope, st)
# for functions that want guessed return types (lazy_node is not None)
# do that now
if (
scope is not None
and scope.ident is not None
and scope.ident.lazy_node is not None
):
if scope.ident.type_obj is not None:
args = scope.ident.type_obj.__args__[0:-1]
ret_type = self._guess_return_value(scope, scope.ident.lazy_node)
# if this function was a property, don't make a callable just set
# the return type
if isinstance(scope.ident.lazy_node, ast.FunctionDef) and any(
[
isinstance(a, ast.Name) and a.id == "property"
for a in scope.ident.lazy_node.decorator_list
]
):
scope.ident.type_obj = ret_type
elif len(args) == 1 and args[0] == ...:
scope.ident.type_obj = Callable[(..., ret_type)]
else:
scope.ident.type_obj = Callable[[a for a in args], ret_type]
scope.ident.lazy_node = None
pass
else:
raise TypeError("Unexpected type of object in pending list")
def _process_module(self):
if self.module is not None:
# start with just the module
modscope = Scope()
modscope.name = "module"
modscope.parsed = self.module
modscope.parent = None
modscope.is_class = False
# set all globals, ignoring reserved ones with __ prefix - so this can be
# globals() without needing extra filtering
for k, v in self.starting_globals.items():
if k.startswith("__"):
continue
id = Ident()
id.line = 0
id.type_obj = v
modscope.set_ident(k, id)
# modules don't have line ranges, so go to the last entry in the body
if len(self.module.body) > 0:
r = _get_linerange(self.module.body[-1])
else:
self.scopes = [modscope]
return
# start with every line pointing to the module scope
self.scopes = [modscope] * (r[1] + 1)
for e in self.module.body:
self._process_stmt(modscope, e)
while len(self.pending) > 0:
self._process_pending(0)