diff --git a/qrenderdoc/Code/pyrenderdoc/parse_reflection.py b/qrenderdoc/Code/pyrenderdoc/parse_reflection.py index c138c6536..21dd9a4f2 100644 --- a/qrenderdoc/Code/pyrenderdoc/parse_reflection.py +++ b/qrenderdoc/Code/pyrenderdoc/parse_reflection.py @@ -199,6 +199,15 @@ def _nopstrings(string: str) -> str: return "".join(text) +class _target_in_gen(ast.AST): + target: ast.AST + gen: ast.comprehension + + def __init__(self, t: ast.AST, g: ast.comprehension): + self.target = t + self.gen = g + + # a given instance of an identifier and its type class Ident: # first line this ident is valid @@ -422,6 +431,44 @@ class PyReflector: return "@@" + err.replace("@", "") + "@@" return Any + def _add_extras_for_gen( + self, + scope: Scope, + tmp_types: Optional[Dict[str, Any]], + generators: List[ast.comprehension], + ) -> Dict[str, Any]: + if tmp_types is None: + extras = {} + else: + extras = tmp_types.copy() + + for g in generators: + iter: Any = self._get_type(scope, g.iter, tmp_types) + + if _is_generic(List, iter): + iter_type = iter.__args__[0] + elif _is_generic(Tuple, iter): + if len(set(iter.__args__)) == 1: + iter_type = iter.__args__[0] + else: + raise TypeError(f"Ambiguous tuple list comp on {iter}") + else: + raise TypeError(f"Unhandled iter list comp on {iter}") + + if isinstance(g.target, ast.Name): + extras[g.target.id] = iter_type + elif ( + isinstance(g.target, ast.Tuple) + and _is_generic(Tuple, iter_type) + and len(g.target.elts) == len(iter_type.__args__) + ): + for i, e in enumerate(g.target.elts): + if not isinstance(e, ast.Name): + raise TypeError(f"Failed unpacking list comp {i} on {e}") + extras[e.id] = iter_type.__args__[i] + + return extras + def _get_type( self, scope: Scope, @@ -470,6 +517,14 @@ class PyReflector: return ident.type_obj return ident + if isinstance(parsed, _target_in_gen): + try: + extras = self._add_extras_for_gen(scope, tmp_types, [parsed.gen]) + except TypeError as err: + return self._type_failure(str(err)) + + return self._get_type(scope, parsed.target, extras) + if isinstance(parsed, ast.Attribute): # detect single-level self.foo and look up in parent if it exists # these identifiers are stored in the parent scope so they're available @@ -804,39 +859,10 @@ class PyReflector: # for list comprehensions / generators we generate a tmp type for the iterator value if isinstance(parsed, ast.ListComp) or isinstance(parsed, ast.GeneratorExp): - if tmp_types is None: - extras = {} - else: - extras = tmp_types.copy() - - for g in parsed.generators: - iter = self._get_type(scope, g.iter, tmp_types) - - if _is_generic(List, iter): - iter_type = iter.__args__[0] - elif _is_generic(Tuple, iter): - if len(set(iter.__args__)) == 1: - iter_type = iter.__args__[0] - else: - return self._type_failure( - f"Ambiguouos tuple list comp on {iter}" - ) - else: - return self._type_failure(f"Unhandle iter list comp on {iter}") - - if isinstance(g.target, ast.Name): - extras[g.target.id] = iter_type - elif ( - isinstance(g.target, ast.Tuple) - and _is_generic(Tuple, iter_type) - and len(g.target.elts) == len(iter_type.__args__) - ): - for i, e in enumerate(g.target.elts): - if not isinstance(e, ast.Name): - return self._type_failure( - f"Failed unpacking list comp {i} on {e}" - ) - extras[e.id] = iter_type.__args__[i] + try: + extras = self._add_extras_for_gen(scope, tmp_types, parsed.generators) + except TypeError as err: + return self._type_failure(str(err)) inner = self._get_type(scope, parsed.elt, extras) if isinstance(inner, str): @@ -1336,7 +1362,6 @@ class PyReflector: multi_fields = [ "orelse", "finalbody", - "generators", "bases", "decorator_list", "targets", @@ -1349,6 +1374,21 @@ class PyReflector: "names", ] + # handle listcomps/generators specially so we can return a hacky thing saying + # which generator to use + if isinstance(parsed, ast.ListComp) or isinstance(parsed, ast.GeneratorExp): + for g in parsed.generators: + ret = self._get_atom_expr(g.target, line, col) + if ret is not None: + return _target_in_gen(ret, g) + ret = self._get_atom_expr(g.iter, line, col) + if ret is not None: + return ret + for ifg in g.ifs: + ret = self._get_atom_expr(ifg, line, col) + if ret is not None: + return ret + if not isinstance(parsed, ast.Lambda): multi_fields += ["body"] @@ -1802,11 +1842,21 @@ if __name__ == "impossible": # ^ # TYPE: List[Tuple[int,str]] unpack_comp = [x for x, y in tuple_list] - # ^ # TYPE: List[int] + # ^ # TYPE: List[int] + # ^ # TYPE: int + # ^ # TYPE: str unpack_comp = [y for x, y in tuple_list] # ^ # TYPE: List[str] + ## Iterators we treat as lists for simplicity + + iter_comp = (x * 2 for x in range(10)) + # ^ # TYPE: List[int] + + iter_comp = next(x * 2 for x in range(10)) + # ^ # TYPE: int + ## subscript/attribute accesses val = list_item[0]