gh-153844: Support AST input in symtable.symtable() (GH-153845)

The builtin compile() accepts an AST object since Python 2.6, but
symtable.symtable() only accepted str and bytes, although the
implementation builds the symbol table from an AST anyway.  Accept an
AST object as well: convert it for the requested compile type, validate
it, honor future statements found in the tree, and build the symbol
table the same way the compiler does.

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: serhiy-storchaka

Doc/library/symtable.rst

@@ -20,6 +20,8 @@ Generating Symbol Tables
 .. function:: symtable(code, filename, compile_type, *, module=None)
 
    Return the toplevel :class:`SymbolTable` for the Python source *code*.
+   *code* can be a string, a bytes object, or an AST object,
+   as for the builtin :func:`compile`.
    *filename* is the name of the file containing the code.  *compile_type* is
    like the *mode* argument to :func:`compile`.
    The optional argument *module* specifies the module name.
@@ -29,6 +31,9 @@ Generating Symbol Tables
    .. versionadded:: 3.15
       Added the *module* parameter.
 
+   .. versionchanged:: next
+      *code* can now be an AST object.
+
 
 Examining Symbol Tables
 -----------------------

Doc/whatsnew/3.16.rst

@@ -379,6 +379,14 @@ shlex
   (Contributed by Jay Berry in :gh:`148846`.)
 
 
+symtable
+--------
+
+* :func:`symtable.symtable` now accepts an AST object,
+  like the builtin :func:`compile`.
+  (Contributed by Serhiy Storchaka in :gh:`153844`.)
+
+
 tkinter
 -------
 

Lib/symtable.py

@@ -20,6 +20,7 @@
 def symtable(code, filename, compile_type, *, module=None):
     """ Return the toplevel *SymbolTable* for the source code.
 
+    *code* can be a string, a bytes object, or an AST object.
     *filename* is the name of the file with the code
     and *compile_type* is the *compile()* mode argument.
     """

Lib/test/test_symtable.py

@@ -2,6 +2,7 @@
 Test the API of the symtable module.
 """
 
+import ast
 import symtable
 import warnings
 import unittest
@@ -554,6 +555,77 @@ def test_loopvar_in_only_one_scope(self):
                 self.assertEqual(len([x for x in ids if x == 'x']), 1)
 
 
+class ASTInputTests(unittest.TestCase):
+    maxDiff = None
+
+    def dump(self, table):
+        return (table.get_name(), table.get_type(), table.get_lineno(),
+                [repr(symbol) for symbol in table.get_symbols()],
+                [self.dump(child) for child in table.get_children()])
+
+    def test_exec(self):
+        top = symtable.symtable(ast.parse(TEST_CODE), "?", "exec")
+        self.assertIsNotNone(find_block(top, "Mine"))
+
+    def test_eval(self):
+        table = symtable.symtable(ast.parse("a + b", mode="eval"), "?", "eval")
+        self.assertEqual(sorted(table.get_identifiers()), ["a", "b"])
+
+    def test_single(self):
+        table = symtable.symtable(ast.parse("x = 1", mode="single"),
+                                  "?", "single")
+        self.assertIn("x", table.get_identifiers())
+
+    def test_same_result_as_string(self):
+        cases = [
+            (TEST_CODE, "exec"),
+            ("from __future__ import annotations\n"
+             "def f(x: int) -> int: return x\n", "exec"),
+            ("[x*y for x in a]", "eval"),
+            ("def f(): pass\n", "single"),
+        ]
+        for source, mode in cases:
+            with self.subTest(source=source, mode=mode):
+                from_str = symtable.symtable(source, "?", mode)
+                from_ast = symtable.symtable(ast.parse(source, mode=mode),
+                                             "?", mode)
+                self.assertEqual(self.dump(from_ast), self.dump(from_str))
+
+    def test_synthesized_ast(self):
+        # An AST created programmatically, without any source.
+        node = ast.Module(body=[
+            ast.FunctionDef(
+                name="f",
+                args=ast.arguments(args=[ast.arg(arg="x")]),
+                body=[ast.Return(ast.Name("x", ast.Load()))])])
+        ast.fix_missing_locations(node)
+        top = symtable.symtable(node, "?", "exec")
+        f = find_block(top, "f")
+        self.assertTrue(f.lookup("x").is_parameter())
+
+    def test_mode_mismatch(self):
+        tree = ast.parse("x = 1")
+        for mode in ("eval", "single"):
+            with self.subTest(mode=mode):
+                with self.assertRaises(TypeError):
+                    symtable.symtable(tree, "?", mode)
+        with self.assertRaises(TypeError):
+            symtable.symtable(ast.parse("x", mode="eval"), "?", "exec")
+
+    def test_invalid_ast(self):
+        node = ast.Expression(ast.Name("x", ast.Store()))
+        ast.fix_missing_locations(node)
+        with self.assertRaises(ValueError):
+            symtable.symtable(node, "?", "eval")
+
+    def test_misplaced_future_import(self):
+        # The parser does not enforce the placement of future imports in
+        # an existing AST; the symbol table construction does.
+        tree = ast.parse("x = 1\nfrom __future__ import annotations\n")
+        with self.assertRaises(SyntaxError):
+            symtable.symtable(tree, "?", "exec")
+
+
 class CommandLineTest(unittest.TestCase):
     maxDiff = None
 

Misc/NEWS.d/next/Library/2026-07-17-14-30-00.gh-issue-153844.pYs7Kq.rst

@@ -0,0 +1,2 @@
+:func:`symtable.symtable` now accepts an AST object,
+like the builtin :func:`compile`.

Modules/clinic/symtablemodule.c.h

@@ -1,4 +1,6 @@
 #include "Python.h"
+#include "pycore_ast.h"           // PyAST_Check()
+#include "pycore_pyarena.h"       // _PyArena_New()
 #include "pycore_pythonrun.h"     // _Py_SourceAsString()
 #include "pycore_symtable.h"      // struct symtable
 
@@ -9,6 +11,29 @@ module _symtable
 /*[clinic end generated code: output=da39a3ee5e6b4b0d input=f4685845a7100605]*/
 
 
+static struct symtable *
+symtable_from_ast(PyObject *source, PyObject *filename, int compile_mode)
+{
+    PyArena *arena = _PyArena_New();
+    if (arena == NULL) {
+        return NULL;
+    }
+    struct symtable *st = NULL;
+    mod_ty mod = PyAST_obj2mod(source, arena, compile_mode);
+    if (mod == NULL || !_PyAST_Validate(mod)) {
+        goto finally;
+    }
+    _PyFutureFeatures future;
+    if (!_PyFuture_FromAST(mod, filename, &future)) {
+        goto finally;
+    }
+    st = _PySymtable_Build(mod, filename, &future);
+finally:
+    _PyArena_Free(arena);
+    return st;
+}
+
+
 /*[clinic input]
 _symtable.symtable
 
@@ -20,37 +45,29 @@ _symtable.symtable
     module as modname: object = None
 
 Return symbol and scope dictionaries used internally by compiler.
+
+The source can be a string, a bytes object, or an AST object.
 [clinic start generated code]*/
 
 static PyObject *
 _symtable_symtable_impl(PyObject *module, PyObject *source,
                         PyObject *filename, const char *startstr,
                         PyObject *modname)
-/*[clinic end generated code: output=235ec5a87a9ce178 input=fbf9adaa33c7070d]*/
+/*[clinic end generated code: output=235ec5a87a9ce178 input=6cadac0485f576a7]*/
 {
     struct symtable *st;
     PyObject *t;
-    int start;
-    PyCompilerFlags cf = _PyCompilerFlags_INIT;
-    PyObject *source_copy = NULL;
-
-    cf.cf_flags = PyCF_SOURCE_IS_UTF8;
-
-    const char *str = _Py_SourceAsString(source, "symtable", "string or bytes", &cf, &source_copy);
-    if (str == NULL) {
-        return NULL;
-    }
+    int compile_mode;
 
     if (strcmp(startstr, "exec") == 0)
-        start = Py_file_input;
+        compile_mode = 0;
     else if (strcmp(startstr, "eval") == 0)
-        start = Py_eval_input;
+        compile_mode = 1;
     else if (strcmp(startstr, "single") == 0)
-        start = Py_single_input;
+        compile_mode = 2;
     else {
         PyErr_SetString(PyExc_ValueError,
            "symtable() arg 3 must be 'exec' or 'eval' or 'single'");
-        Py_XDECREF(source_copy);
         return NULL;
     }
     if (modname == Py_None) {
@@ -60,11 +77,30 @@ _symtable_symtable_impl(PyObject *module, PyObject *source,
         PyErr_Format(PyExc_TypeError,
                      "symtable() argument 'module' must be str or None, not %T",
                      modname);
-        Py_XDECREF(source_copy);
         return NULL;
     }
-    st = _Py_SymtableStringObjectFlags(str, filename, start, &cf, modname);
-    Py_XDECREF(source_copy);
+
+    if (PyAST_Check(source)) {
+        st = symtable_from_ast(source, filename, compile_mode);
+    }
+    else {
+        static const int starts[] = {
+            Py_file_input, Py_eval_input, Py_single_input};
+        PyCompilerFlags cf = _PyCompilerFlags_INIT;
+        PyObject *source_copy = NULL;
+
+        cf.cf_flags = PyCF_SOURCE_IS_UTF8;
+        const char *str = _Py_SourceAsString(source, "symtable",
+                                             "string, bytes or AST",
+                                             &cf, &source_copy);
+        if (str == NULL) {
+            return NULL;
+        }
+        st = _Py_SymtableStringObjectFlags(str, filename,
+                                           starts[compile_mode], &cf,
+                                           modname);
+        Py_XDECREF(source_copy);
+    }
     if (st == NULL) {
         return NULL;
     }