From 755dab719dfc924dd8aef46f67512dabb8f25071 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?B=C3=A9n=C3=A9dikt=20Tran?=
 <10796600+picnixz@users.noreply.github.com>
Date: Wed, 12 Jun 2024 13:14:50 +0200
Subject: [PATCH] gh-120029: make `symtable.Symbol.__repr__` correctly reflect
 the compiler's flags, add methods  (#120099)

Expose :class:`symtable.Symbol` methods :meth:`~symtable.Symbol.is_free_class`,
:meth:`~symtable.Symbol.is_comp_iter` and :meth:`~symtable.Symbol.is_comp_cell`.

---------

Co-authored-by: Carl Meyer <carl@oddbird.net>
---
 Doc/library/symtable.rst                      | 34 +++++++++++++++++++
 Doc/whatsnew/3.14.rst                         | 11 ++++++
 Include/internal/pycore_symtable.h            |  2 +-
 Lib/symtable.py                               | 32 ++++++++++++++---
 Lib/test/test_symtable.py                     | 21 ++++++++++++
 ...-06-05-11-03-10.gh-issue-120029.QBsw47.rst |  4 +++
 Modules/symtablemodule.c                      |  2 ++
 7 files changed, 100 insertions(+), 6 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Library/2024-06-05-11-03-10.gh-issue-120029.QBsw47.rst

diff --git a/Doc/library/symtable.rst b/Doc/library/symtable.rst
index e17a33f7feb1ab..050a941d9d0516 100644
--- a/Doc/library/symtable.rst
+++ b/Doc/library/symtable.rst
@@ -155,6 +155,8 @@ Examining Symbol Tables
 
       Return ``True`` if the symbol is a type parameter.
 
+      .. versionadded:: 3.14
+
    .. method:: is_global()
 
       Return ``True`` if the symbol is global.
@@ -182,10 +184,42 @@ Examining Symbol Tables
       Return ``True`` if the symbol is referenced in its block, but not assigned
       to.
 
+   .. method:: is_free_class()
+
+      Return *True* if a class-scoped symbol is free from
+      the perspective of a method.
+
+      Consider the following example::
+
+         def f():
+             x = 1  # function-scoped
+             class C:
+                 x = 2  # class-scoped
+                 def method(self):
+                     return x
+
+      In this example, the class-scoped symbol ``x`` is considered to
+      be free from the perspective of ``C.method``, thereby allowing
+      the latter to return *1* at runtime and not *2*.
+
+      .. versionadded:: 3.14
+
    .. method:: is_assigned()
 
       Return ``True`` if the symbol is assigned to in its block.
 
+   .. method:: is_comp_iter()
+
+      Return ``True`` if the symbol is a comprehension iteration variable.
+
+      .. versionadded:: 3.14
+
+   .. method:: is_comp_cell()
+
+      Return ``True`` if the symbol is a cell in an inlined comprehension.
+
+      .. versionadded:: 3.14
+
    .. method:: is_namespace()
 
       Return ``True`` if name binding introduces new namespace.
diff --git a/Doc/whatsnew/3.14.rst b/Doc/whatsnew/3.14.rst
index b77ff30a8fbbee..b357553735e8bb 100644
--- a/Doc/whatsnew/3.14.rst
+++ b/Doc/whatsnew/3.14.rst
@@ -100,6 +100,17 @@ os
   by :func:`os.unsetenv`, or made outside Python in the same process.
   (Contributed by Victor Stinner in :gh:`120057`.)
 
+symtable
+--------
+
+* Expose the following :class:`symtable.Symbol` methods:
+
+  * :meth:`~symtable.Symbol.is_free_class`
+  * :meth:`~symtable.Symbol.is_comp_iter`
+  * :meth:`~symtable.Symbol.is_comp_cell`
+
+  (Contributed by Bénédikt Tran in :gh:`120029`.)
+
 
 Optimizations
 =============
diff --git a/Include/internal/pycore_symtable.h b/Include/internal/pycore_symtable.h
index 5d544765237df5..1be48edc80c830 100644
--- a/Include/internal/pycore_symtable.h
+++ b/Include/internal/pycore_symtable.h
@@ -154,7 +154,7 @@ extern PyObject* _Py_Mangle(PyObject *p, PyObject *name);
 #define DEF_BOUND (DEF_LOCAL | DEF_PARAM | DEF_IMPORT)
 
 /* GLOBAL_EXPLICIT and GLOBAL_IMPLICIT are used internally by the symbol
-   table.  GLOBAL is returned from PyST_GetScope() for either of them.
+   table.  GLOBAL is returned from _PyST_GetScope() for either of them.
    It is stored in ste_symbols at bits 13-16.
 */
 #define SCOPE_OFFSET 12
diff --git a/Lib/symtable.py b/Lib/symtable.py
index af65e93e68eda4..d6ac1f527ba8ba 100644
--- a/Lib/symtable.py
+++ b/Lib/symtable.py
@@ -4,7 +4,10 @@
 from _symtable import (
     USE,
     DEF_GLOBAL, DEF_NONLOCAL, DEF_LOCAL,
-    DEF_PARAM, DEF_TYPE_PARAM, DEF_IMPORT, DEF_BOUND, DEF_ANNOT,
+    DEF_PARAM, DEF_TYPE_PARAM,
+    DEF_FREE_CLASS,
+    DEF_IMPORT, DEF_BOUND, DEF_ANNOT,
+    DEF_COMP_ITER, DEF_COMP_CELL,
     SCOPE_OFF, SCOPE_MASK,
     FREE, LOCAL, GLOBAL_IMPLICIT, GLOBAL_EXPLICIT, CELL
 )
@@ -158,6 +161,10 @@ def get_children(self):
                 for st in self._table.children]
 
 
+def _get_scope(flags):  # like _PyST_GetScope()
+    return (flags >> SCOPE_OFF) & SCOPE_MASK
+
+
 class Function(SymbolTable):
 
     # Default values for instance variables
@@ -183,7 +190,7 @@ def get_locals(self):
         """
         if self.__locals is None:
             locs = (LOCAL, CELL)
-            test = lambda x: ((x >> SCOPE_OFF) & SCOPE_MASK) in locs
+            test = lambda x: _get_scope(x) in locs
             self.__locals = self.__idents_matching(test)
         return self.__locals
 
@@ -192,7 +199,7 @@ def get_globals(self):
         """
         if self.__globals is None:
             glob = (GLOBAL_IMPLICIT, GLOBAL_EXPLICIT)
-            test = lambda x:((x >> SCOPE_OFF) & SCOPE_MASK) in glob
+            test = lambda x: _get_scope(x) in glob
             self.__globals = self.__idents_matching(test)
         return self.__globals
 
@@ -207,7 +214,7 @@ def get_frees(self):
         """Return a tuple of free variables in the function.
         """
         if self.__frees is None:
-            is_free = lambda x:((x >> SCOPE_OFF) & SCOPE_MASK) == FREE
+            is_free = lambda x: _get_scope(x) == FREE
             self.__frees = self.__idents_matching(is_free)
         return self.__frees
 
@@ -234,7 +241,7 @@ class Symbol:
     def __init__(self, name, flags, namespaces=None, *, module_scope=False):
         self.__name = name
         self.__flags = flags
-        self.__scope = (flags >> SCOPE_OFF) & SCOPE_MASK # like PyST_GetScope()
+        self.__scope = _get_scope(flags)
         self.__namespaces = namespaces or ()
         self.__module_scope = module_scope
 
@@ -303,6 +310,11 @@ def is_free(self):
         """
         return bool(self.__scope == FREE)
 
+    def is_free_class(self):
+        """Return *True* if a class-scoped symbol is free from
+        the perspective of a method."""
+        return bool(self.__flags & DEF_FREE_CLASS)
+
     def is_imported(self):
         """Return *True* if the symbol is created from
         an import statement.
@@ -313,6 +325,16 @@ def is_assigned(self):
         """Return *True* if a symbol is assigned to."""
         return bool(self.__flags & DEF_LOCAL)
 
+    def is_comp_iter(self):
+        """Return *True* if the symbol is a comprehension iteration variable.
+        """
+        return bool(self.__flags & DEF_COMP_ITER)
+
+    def is_comp_cell(self):
+        """Return *True* if the symbol is a cell in an inlined comprehension.
+        """
+        return bool(self.__flags & DEF_COMP_CELL)
+
     def is_namespace(self):
         """Returns *True* if name binding introduces new namespace.
 
diff --git a/Lib/test/test_symtable.py b/Lib/test/test_symtable.py
index a4b111e865c86e..903c6d66f50964 100644
--- a/Lib/test/test_symtable.py
+++ b/Lib/test/test_symtable.py
@@ -304,6 +304,27 @@ def test_symbol_repr(self):
         self.assertEqual(repr(self.GenericMine.lookup("T")),
                          "<symbol 'T': LOCAL, DEF_LOCAL|DEF_TYPE_PARAM>")
 
+        st1 = symtable.symtable("[x for x in [1]]", "?", "exec")
+        self.assertEqual(repr(st1.lookup("x")),
+                         "<symbol 'x': LOCAL, USE|DEF_LOCAL|DEF_COMP_ITER>")
+
+        st2 = symtable.symtable("[(lambda: x) for x in [1]]", "?", "exec")
+        self.assertEqual(repr(st2.lookup("x")),
+                         "<symbol 'x': CELL, DEF_LOCAL|DEF_COMP_ITER|DEF_COMP_CELL>")
+
+        st3 = symtable.symtable("def f():\n"
+                                "   x = 1\n"
+                                "   class A:\n"
+                                "       x = 2\n"
+                                "       def method():\n"
+                                "           return x\n",
+                                "?", "exec")
+        # child 0 is for __annotate__
+        func_f = st3.get_children()[1]
+        class_A = func_f.get_children()[0]
+        self.assertEqual(repr(class_A.lookup('x')),
+                         "<symbol 'x': LOCAL, DEF_LOCAL|DEF_FREE_CLASS>")
+
     def test_symtable_entry_repr(self):
         expected = f"<symtable entry top({self.top.get_id()}), line {self.top.get_lineno()}>"
         self.assertEqual(repr(self.top._table), expected)
diff --git a/Misc/NEWS.d/next/Library/2024-06-05-11-03-10.gh-issue-120029.QBsw47.rst b/Misc/NEWS.d/next/Library/2024-06-05-11-03-10.gh-issue-120029.QBsw47.rst
new file mode 100644
index 00000000000000..d1b2c592a113ce
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2024-06-05-11-03-10.gh-issue-120029.QBsw47.rst
@@ -0,0 +1,4 @@
+Expose :class:`symtable.Symbol` methods :meth:`~symtable.Symbol.is_free_class`,
+:meth:`~symtable.Symbol.is_comp_iter` and :meth:`~symtable.Symbol.is_comp_cell`.
+Patch by Bénédikt Tran.
+
diff --git a/Modules/symtablemodule.c b/Modules/symtablemodule.c
index 63c4dd4225298d..b39b59bf7b06bf 100644
--- a/Modules/symtablemodule.c
+++ b/Modules/symtablemodule.c
@@ -81,6 +81,8 @@ symtable_init_constants(PyObject *m)
     if (PyModule_AddIntMacro(m, DEF_IMPORT) < 0) return -1;
     if (PyModule_AddIntMacro(m, DEF_BOUND) < 0) return -1;
     if (PyModule_AddIntMacro(m, DEF_ANNOT) < 0) return -1;
+    if (PyModule_AddIntMacro(m, DEF_COMP_ITER) < 0) return -1;
+    if (PyModule_AddIntMacro(m, DEF_COMP_CELL) < 0) return -1;
 
     if (PyModule_AddIntConstant(m, "TYPE_FUNCTION", FunctionBlock) < 0)
         return -1;