diff --git a/pglast/ast.py b/pglast/ast.py index bd2c0c9..5834fea 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 12beb3d..ff08b40 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 0832d9e..88f5159 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,)