gh-157056: Fix STRING format for comprehension and lambda annotations

Lib/annotationlib.py

@@ -750,13 +750,22 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
             argdefs=annotate.__defaults__,
             kwdefaults=annotate.__kwdefaults__,
         )
-        annos = func(Format.VALUE_WITH_FAKE_GLOBALS)
+        try:
+            annos = func(Format.VALUE_WITH_FAKE_GLOBALS)
+        except ValueError as exc:
+            # Dict comprehensions such as `{k: v for k, v in items}` unpack
+            # each iterated element. Fake-globals iteration yields a single
+            # starred stringifier, so unpacking raises ValueError. Recover
+            # the original annotation text from source when we can.
+            # Changing _Stringifier.__iter__ would break [*a] stringification.
+            if _is_unpack_value_error(exc):
+                sourced = _string_annotations_from_source(owner)
+                if sourced is not None:
+                    return sourced
+            raise
         if _is_evaluate:
             return _stringify_single(annos)
-        return {
-            key: _stringify_single(val)
-            for key, val in annos.items()
-        }
+        return _stringify_annotation_dict(annos, owner)
     elif format == Format.FORWARDREF:
         # FORWARDREF is implemented similarly to STRING, but there are two changes,
         # at the beginning and the end of the process.
@@ -878,6 +887,175 @@ def _build_closure(annotate, owner, is_class, stringifier_dict, *, allow_evaluat
     return tuple(new_closure), cell_dict
 
 
+def _string_annotations_from_source(obj):
+    """Best-effort STRING annotations reconstructed from *obj*'s source AST.
+
+    Used when fake-globals evaluation cannot stringify an annotation (dict
+    comprehensions, lambdas, generator expressions). Returns None if source
+    is unavailable. inspect is imported lazily because it imports this module.
+    """
+    if obj is None:
+        return None
+    try:
+        import inspect
+        import textwrap
+        source = inspect.getsource(obj)
+    except (OSError, TypeError, RecursionError):
+        return None
+    source = textwrap.dedent(source)
+    try:
+        tree = ast.parse(source)
+    except SyntaxError:
+        return None
+    if not tree.body:
+        return None
+    node = tree.body[0]
+    result = {}
+    if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+        for arg in (
+            *node.args.posonlyargs,
+            *node.args.args,
+            *node.args.kwonlyargs,
+        ):
+            if arg.annotation is not None:
+                result[arg.arg] = _annotation_ast_to_string(arg.annotation)
+        if node.args.vararg is not None and node.args.vararg.annotation is not None:
+            result[node.args.vararg.arg] = _annotation_ast_to_string(
+                node.args.vararg.annotation
+            )
+        if node.args.kwarg is not None and node.args.kwarg.annotation is not None:
+            result[node.args.kwarg.arg] = _annotation_ast_to_string(
+                node.args.kwarg.annotation
+            )
+        if node.returns is not None:
+            result["return"] = _annotation_ast_to_string(node.returns)
+        return result
+    if isinstance(node, ast.ClassDef):
+        _collect_class_annassigns(node.body, result)
+        return result
+    return None
+
+
+def _annotation_ast_to_string(node):
+    # Match fake-globals STRING output for quoted annotations: a: "int" -> "int".
+    if isinstance(node, ast.Constant) and isinstance(node.value, str):
+        return node.value
+    return ast.unparse(node)
+
+
+def _collect_class_annassigns(body, result):
+    """Collect class annotations, including those inside simple control flow.
+
+    Nested functions and classes are skipped so their annotations are not
+    attributed to the enclosing class.
+    """
+    for stmt in body:
+        if isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name):
+            result[stmt.target.id] = _annotation_ast_to_string(stmt.annotation)
+        elif isinstance(stmt, ast.If):
+            _collect_class_annassigns(stmt.body, result)
+            _collect_class_annassigns(stmt.orelse, result)
+        elif isinstance(stmt, ast.Try):
+            _collect_class_annassigns(stmt.body, result)
+            for handler in stmt.handlers:
+                _collect_class_annassigns(handler.body, result)
+            _collect_class_annassigns(stmt.orelse, result)
+            _collect_class_annassigns(stmt.finalbody, result)
+        elif isinstance(stmt, (ast.With, ast.AsyncWith)):
+            _collect_class_annassigns(stmt.body, result)
+        elif isinstance(stmt, (ast.For, ast.AsyncFor, ast.While)):
+            _collect_class_annassigns(stmt.body, result)
+            _collect_class_annassigns(stmt.orelse, result)
+        elif isinstance(stmt, ast.Match):
+            for case in stmt.cases:
+                _collect_class_annassigns(case.body, result)
+
+
+def _is_unpack_value_error(exc):
+    if not exc.args:
+        return False
+    msg = exc.args[0]
+    return isinstance(msg, str) and "values to unpack" in msg
+
+
+def _is_runtime_constructed(value):
+    """True if *value* was created at annotation-eval time and has no AST.
+
+    Lambdas and generator expressions are syntax, not name lookups, so the
+    fake-globals stringifier never sees them. Their repr() embeds a memory
+    address and is not a valid annotation string.
+    """
+    return isinstance(value, (
+        types.FunctionType,
+        types.BuiltinFunctionType,
+        types.MethodType,
+        types.GeneratorType,
+        types.AsyncGeneratorType,
+        types.CoroutineType,
+    ))
+
+
+def _contains_runtime_constructed(value):
+    if _is_runtime_constructed(value):
+        return True
+    if isinstance(value, dict):
+        return any(
+            _contains_runtime_constructed(k) or _contains_runtime_constructed(v)
+            for k, v in value.items()
+        )
+    if isinstance(value, (list, tuple, set, frozenset)):
+        return any(_contains_runtime_constructed(v) for v in value)
+    return False
+
+
+def _stringify_annotation_dict(annos, owner):
+    # Only read source when fake-globals produced a runtime object we cannot
+    # stringify (lambda, genexpr, or a container holding one), or when the
+    # fallback text still embeds a memory address.
+    sourced = _sentinel
+    result = {}
+    for key, val in annos.items():
+        text = _stringify_single(val)
+        needs_source = (
+            _contains_runtime_constructed(val)
+            or "0x" in text.lower()
+        )
+        if needs_source:
+            if sourced is _sentinel:
+                sourced = _string_annotations_from_source(owner)
+            if sourced is not None and key in sourced:
+                result[key] = sourced[key]
+                continue
+        result[key] = text
+    return result
+
+
+def _stringify_container(anno):
+    if isinstance(anno, list):
+        inner = ", ".join(_stringify_single(v) for v in anno)
+        return f"[{inner}]"
+    if isinstance(anno, tuple):
+        inner = ", ".join(_stringify_single(v) for v in anno)
+        if len(anno) == 1:
+            inner += ","
+        return f"({inner})"
+    if isinstance(anno, set):
+        if not anno:
+            return "set()"
+        inner = ", ".join(_stringify_single(v) for v in anno)
+        return f"{{{inner}}}"
+    if isinstance(anno, frozenset):
+        if not anno:
+            return "frozenset()"
+        inner = ", ".join(_stringify_single(v) for v in anno)
+        return f"frozenset({{{inner}}})"
+    inner = ", ".join(
+        f"{_stringify_single(k)}: {_stringify_single(v)}"
+        for k, v in anno.items()
+    )
+    return f"{{{inner}}}"
+
+
 def _stringify_single(anno):
     if anno is ...:
         return "..."
@@ -886,6 +1064,15 @@ def _stringify_single(anno):
         return anno
     elif isinstance(anno, _Template):
         return ast.unparse(_template_to_ast(anno))
+    elif _is_runtime_constructed(anno):
+        # Lambdas and generator expressions are syntax, not name lookups.
+        # repr() embeds a memory address; type_repr() is stable.
+        return type_repr(anno)
+    elif (
+        isinstance(anno, (list, tuple, set, frozenset, dict))
+        and _contains_runtime_constructed(anno)
+    ):
+        return _stringify_container(anno)
     else:
         return repr(anno)
 
@@ -1089,10 +1276,25 @@ def type_repr(value):
     repr() for most objects.
 
     """
-    if isinstance(value, (type, types.FunctionType, types.BuiltinFunctionType)):
-        if value.__module__ == "builtins":
+    if isinstance(value, (
+        type,
+        types.FunctionType,
+        types.BuiltinFunctionType,
+        types.MethodType,
+    )):
+        if getattr(value, "__module__", None) == "builtins":
             return value.__qualname__
-        return f"{value.__module__}.{value.__qualname__}"
+        module = getattr(value, "__module__", None)
+        if module:
+            return f"{module}.{value.__qualname__}"
+        return value.__qualname__
+    elif isinstance(value, (
+        types.GeneratorType,
+        types.AsyncGeneratorType,
+        types.CoroutineType,
+    )):
+        # repr() of these objects embeds a memory address.
+        return value.__qualname__
     elif isinstance(value, _Template):
         tree = _template_to_ast(value)
         return ast.unparse(tree)