From 901c75e6a8671db8deeb004ce368bf37b76f96e3 Mon Sep 17 00:00:00 2001
From: Lele Gaifax <lele@metapensiero.it>
Date: Sat, 28 May 2022 14:25:04 +0200
Subject: [PATCH] Add an "ancestors" slot to all AST nodes

This is the first step to implement issue #80, making it easier to get
rid of the old Node wrappers, whose main functionality in v3 was to keep
track of the "parent_node"/"parent_member" information.
---
 pglast/ast.py        | 4 ++--
 pglast/stream.py     | 8 +++++++-
 tools/extract_ast.py | 4 ++--
 3 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/pglast/ast.py b/pglast/ast.py
index bd2c0c9d..5834fea3 100644
--- a/pglast/ast.py
+++ b/pglast/ast.py
@@ -58,7 +58,7 @@ def __repr__(self):
 class Node:
     "Base class for all AST nodes."
 
-    __slots__ = ()
+    __slots__ = ('ancestors',)
 
     def __init__(self, data):
         if not isinstance(data, dict):  # pragma: no cover
@@ -141,7 +141,7 @@ def __setattr__(self, name, value):
         attribute, raising opportune exception when that is not possible.
         '''
 
-        if value is not None:
+        if value is not None and name in self.__slots__:
             ctype, ptype, adaptor = self.__slots__[name]
             if not isinstance(ptype, tuple):
                 ptype = (ptype,)
diff --git a/pglast/stream.py b/pglast/stream.py
index 12beb3df..ff08b404 100644
--- a/pglast/stream.py
+++ b/pglast/stream.py
@@ -11,8 +11,8 @@
 from re import match
 from sys import stderr
 
-from . import parse_plpgsql, parse_sql
 from .node import List, Missing, Node, Scalar
+from . import ast, parse_plpgsql, parse_sql, visitors
 from .keywords import RESERVED_KEYWORDS, TYPE_FUNC_NAME_KEYWORDS
 from .printers import get_printer_for_node_tag, get_special_function
 
@@ -176,6 +176,12 @@ def __call__(self, sql, plpgsql=False):
                              " a node.Node instance, a node.List, an ast.Node or tuple of"
                              " them, got %r" % type(sql))
 
+        class UpdateAncestors(visitors.Visitor):
+            def visit(self, ancestors, node):
+                node.ancestors = ancestors
+
+        UpdateAncestors()(sql)
+
         first = True
         for statement in sql:
             if first:
diff --git a/tools/extract_ast.py b/tools/extract_ast.py
index 0832d9ee..88f5159a 100644
--- a/tools/extract_ast.py
+++ b/tools/extract_ast.py
@@ -78,7 +78,7 @@ def __repr__(self):
 class Node:
     "Base class for all AST nodes."
 
-    __slots__ = ()
+    __slots__ = ('ancestors',)
 
     def __init__(self, data):
         if not isinstance(data, dict):  # pragma: no cover
@@ -161,7 +161,7 @@ def __setattr__(self, name, value):
         attribute, raising opportune exception when that is not possible.
         '''
 
-        if value is not None:
+        if value is not None and name in self.__slots__:
             ctype, ptype, adaptor = self.__slots__[name]
             if not isinstance(ptype, tuple):
                 ptype = (ptype,)