Repository: kayak/pypika
Pull request #364 (2019-12-27 10:37:13)

Added nodes_ and find_ functions to search the object graph
Replaced some utility functions which look for entities inside a pypika query object graph with a more generalized nodes_ and find which filters object nodes based on object type


Patch:
diff --git a/pypika/dialects.py b/pypika/dialects.py
index 5afab40..fec7fee 100644
--- a/pypika/dialects.py
+++ b/pypika/dialects.py
@@ -377,7 +377,7 @@ class PostgreQueryBuilder(QueryBuilder):
                 self._return_other(self.wrap_constant(term, self._wrapper_cls))
 
     def _validate_returning_term(self, term):
-        for field in term.fields():
+        for field in term.fields_():
             if not any([self._insert_table, self._update_table, self._delete_from]):
                 raise QueryException("Returning can't be used in this query")
             if (
diff --git a/pypika/queries.py b/pypika/queries.py
index 17cffef..d469e86 100644
--- a/pypika/queries.py
+++ b/pypika/queries.py
@@ -11,6 +11,7 @@ from pypika.terms import (
     Field,
     Function,
     Index,
+    Node,
     Rollup,
     Star,
     Term,
@@ -32,7 +33,7 @@ __author__ = "Timothy Heys"
 __email__ = "theys@kayak.com"
 
 
-class Selectable:
+class Selectable(Node):
     def __init__(self, alias):
         self.alias = alias
 
@@ -55,6 +56,9 @@ class Selectable:
     def __getitem__(self, name):
         return self.field(name)
 
+    def get_table_name(self):
+        return self.alias
+
 
 class AliasedQuery(Selectable):
     def __init__(self, name, query=None):
@@ -132,6 +136,9 @@ class Table(Selectable):
         self._table_name = name
         self._schema = self._init_schema(schema)
 
+    def get_table_name(self):
+        return self.alias or self._table_name
+
     def get_sql(self, **kwargs):
         quote_char = kwargs.get("quote_char")
         # FIXME escape
@@ -838,6 +845,8 @@ class QueryBuilder(Selectable, Term):
             return Joiner(self, item, how, type_label="table")
 
         elif isinstance(item, QueryBuilder):
+            if item.alias is None:
+                self._tag_subquery(item)
             return Joiner(self, item, how, type_label="subquery")
 
         elif isinstance(item, AliasedQuery):
@@ -889,12 +898,16 @@ class QueryBuilder(Selectable, Term):
     def __mul__(self, other):
         return self.union_all(other)
 
+    @builder
+    def slice(self, slice):
+        self._offset = slice.start
+        self._limit = slice.stop
+
     @builder
     def __getitem__(self, item):
         if not isinstance(item, slice):
-            raise TypeError("Query' object is not subscriptable")
-        self._offset = item.start
-        self._limit = item.stop
+            return super().__getitem__(item)
+        return self.slice(item)
 
     @staticmethod
     def _list_aliases(field_set, quote_char=None):
@@ -937,18 +950,14 @@ class QueryBuilder(Selectable, Term):
     def _select_other(self, function):
         self._selects.append(function)
 
-    def fields(self):
+    def fields_(self):
         # Don't return anything here. Subqueries have their own fields.
         return []
 
     def do_join(self, join):
         base_tables = self._from + [self._update_table] + self._with
-
         join.validate(base_tables, self._joins)
 
-        if isinstance(join.item, QueryBuilder) and join.item.alias is None:
-            self._tag_subquery(join.item)
-
         table_in_query = any(
             isinstance(clause, Table) and join.item in base_tables
             for clause in base_tables
@@ -970,14 +979,16 @@ class QueryBuilder(Selectable, Term):
         """
         base_tables = self._from + [self._update_table]
 
-        for field in term.fields():
+        for field in term.fields_():
             table_in_base_tables = field.table in base_tables
             table_in_joins = field.table in [join.item for join in self._joins]
-            if (
-                field.table is not None
-                and not table_in_base_tables
-                and not table_in_joins
-                and field.table != self._update_table
+            if all(
+                [
+                    field.table is not None,
+                    not table_in_base_tables,
+                    not table_in_joins,
+                    field.table != self._update_table,
+                ]
             ):
                 return False
         return True
@@ -1431,7 +1442,7 @@ class JoinOn(Join):
         )
 
     def validate(self, _from, _joins):
-        criterion_tables = set([f.table for f in self.criterion.fields()])
+        criterion_tables = set([f.table for f in self.criterion.fields_()])
         available_tables = set(_from) | {join.item for join in _joins} | {self.item}
         missing_tables = criterion_tables - available_tables
         if missing_tables:
diff --git a/pypika/terms.py b/pypika/terms.py
index 4d68cdf..9b36341 100644
--- a/pypika/terms.py
+++ b/pypika/terms.py
@@ -1,9 +1,7 @@
 import inspect
-import itertools
 import re
 from datetime import date
 from enum import Enum
-
 from typing import (
     Iterable,
     Union,
@@ -36,7 +34,15 @@ __author__ = "Timothy Heys"
 __email__ = "theys@kayak.com"
 
 
-class Term:
+class Node:
+    def nodes_(self):
+        yield self
+
+    def find_(self, type):
+        return [node for node in self.nodes_() if isinstance(node, type)]
+
+
+class Term(Node):
     is_aggregate = False
 
     def __init__(self, alias=None):
@@ -48,7 +54,12 @@ class Term:
 
     @property
     def tables_(self):
-        return set()
+        from pypika import Table
+
+        return set(self.find_(Table))
+
+    def fields_(self):
+        return set(self.find_(Field))
 
     @staticmethod
     def wrap_constant(val, wrapper_cls=None):
@@ -109,9 +120,6 @@ class Term:
         """
         return self
 
-    def fields(self):
-        return [self]
-
     def eq(self, other):
         return self == other
 
@@ -254,9 +262,6 @@ class Parameter(Term):
         super(Parameter, self).__init__()
         self.placeholder = placeholder
 
-    def fields(self):
-        return []
-
     def get_sql(self, **kwargs):
         return str(self.placeholder)
 
@@ -281,9 +286,6 @@ class ValueWrapper(Term):
         super(ValueWrapper, self).__init__(alias)
         self.value = value
 
-    def fields(self):
-        return []
-
     def get_value_sql(self, **kwargs):
         quote_char = kwargs.get("secondary_quote_char") or ""
 
@@ -398,9 +400,6 @@ class Values(Term):
 
 
 class NullValue(Term):
-    def fields(self):
-        return []
-
     def get_sql(self, **kwargs):
         sql = "NULL"
         return format_alias_sql(sql, self.alias, **kwargs)
@@ -434,9 +433,6 @@ class Criterion(Term):
 
         return crit
 
-    def fields(self):
-        raise NotImplementedError()
-
     def get_sql(self):
         raise NotImplementedError()
 
@@ -461,12 +457,10 @@ class Field(Criterion, JSON):
         self.name = name
         self.table = table
 
-    def fields(self):
-        return [self]
-
-    @property
-    def tables_(self):
-        return {self.table}
+    def nodes_(self):
+        yield self
+        if self.table is not None:
+            yield from self.table.nodes_()
 
     @builder
     def replace_table(self, current_table, new_table):
@@ -494,11 +488,9 @@ class Field(Criterion, JSON):
 
         # Need to add namespace if the table has an alias
         if self.table and (with_namespace or self.table.alias):
+            table_name = self.table.get_table_name()
             field_sql = "{namespace}.{name}".format(
-                namespace=format_quotes(
-                    self.table.alias or self.table._table_name, quote_char
-                ),
-                name=field_sql,
+                namespace=format_quotes(table_name, quote_char), name=field_sql,
             )
 
         field_alias = getattr(self, "alias", None)
@@ -523,11 +515,10 @@ class Star(Field):
     def __init__(self, table=None):
         super(Star, self).__init__("*", table=table)
 
-    @property
-    def tables_(self):
-        if self.table is None:
-            return {}
-        return {self.table}
+    def nodes_(self):
+        yield self
+        if self.table is not None:
+            yield from self.table.nodes_()
 
     def get_sql(
         self, with_alias=False, with_namespace=False, quote_char=None, **kwargs
@@ -544,8 +535,10 @@ class Tuple(Criterion):
         super(Tuple, self).__init__()
         self.values = [self.wrap_constant(value) for value in values]
 
-    def fields(self):
-        return list(itertools.chain(*[value.fields() for value in self.values]))
+    def nodes_(self):
+        yield self
+        for value in self.values:
+            yield from value.nodes_()
 
     def get_sql(self, **kwargs):
         return "({})".format(",".join(term.get_sql(**kwargs) for term in self.values))
@@ -602,12 +595,11 @@ class NestedCriterion(Criterion):
         self.right = right
         self.nested = nested
 
-    def fields(self):
-        return list(
-            itertools.chain(
-                self.right.fields(), self.left.fields(), self.nested.fields()
-            )
-        )
+    def nodes_(self):
+        yield self
+        yield from self.right.nodes_()
+        yield from self.left.nodes_()
+        yield from self.nested.nodes_()
 
     @property
     def is_aggregate(self):
@@ -615,10 +607,6 @@ class NestedCriterion(Criterion):
             [term.is_aggregate for term in [self.left, self.right, self.nested]]
         )
 
-    @property
-    def tables_(self):
-        return self.left.tables_ | self.right.tables_ | self.nested.tables_
-
     @builder
     def replace_table(self, current_table, new_table):
         """
@@ -670,16 +658,17 @@ class BasicCriterion(Criterion):
         self.left = left
         self.right = right
 
+    def nodes_(self):
+        yield self
+        yield from self.right.nodes_()
+        yield from self.left.nodes_()
+
     @property
     def is_aggregate(self):
         return resolve_is_aggregate(
             [term.is_aggregate for term in [self.left, self.right]]
         )
 
-    @property
-    def tables_(self):
-        return self.left.tables_ | self.right.tables_
-
     @builder
     def replace_table(self, current_table, new_table):
         """
@@ -695,9 +684,6 @@ class BasicCriterion(Criterion):
         self.left = self.left.replace_table(current_table, new_table)
         self.right = self.right.replace_table(current_table, new_table)
 
-    def fields(self):
-        return self.left.fields() + self.right.fields()
-
     def get_sql(self, quote_char='"', with_alias=False, **kwargs):
         sql = "{left}{comparator}{right}".format(
             comparator=self.comparator.value,
@@ -727,9 +713,10 @@ class ContainsCriterion(Criterion):
         self.container = container
         self._is_negated = False
 
-    @property
-    def tables_(self):
-        return self.term.tables_
+    def nodes_(self):
+        yield self
+        yield from self.term.nodes_()
+        yield from self.container.nodes_()
 
     @property
     def is_aggregate(self):
@@ -749,9 +736,6 @@ class ContainsCriterion(Criterion):
         """
         self.term = self.term.replace_table(current_table, new_table)
 
-    def fields(self):
-        return self.term.fields() if self.term.fields else []
-
     def get_sql(self, subquery=None, **kwargs):
         return "{term} {not_}IN {container}".format(
             term=self.term.get_sql(**kwargs),
@@ -759,9 +743,9 @@ class ContainsCriterion(Criterion):
             not_="NOT " if self._is_negated else "",
         )
 
+    @builder
     def negate(self):
         self._is_negated = True
-        return self
 
 
 class BetweenCriterion(Criterion):
@@ -771,9 +755,11 @@ class BetweenCriterion(Criterion):
         self.start = start
         self.end = end
 
-    @property
-    def tables_(self):
-        return self.term.tables_
+    def nodes_(self):
+        yield self
+        yield from self.term.nodes_()
+        yield from self.start.nodes_()
+        yield from self.end.nodes_()
 
     @property
     def is_aggregate(self):
@@ -801,9 +787,6 @@ class BetweenCriterion(Criterion):
             end=self.end.get_sql(**kwargs),
         )
 
-    def fields(self):
-        return self.term.fields() if self.term.fields else []
-
 
 class BitwiseAndCriterion(Criterion):
     def __init__(self, term, value, alias=None):
@@ -811,9 +794,10 @@ class BitwiseAndCriterion(Criterion):
         self.term = term
         self.value = value
 
-    @property
-    def tables_(self):
-        return self.term.tables_
+    def nodes_(self):
+        yield self
+        yield from self.term.nodes_()
+        yield from self.value.nodes_()
 
     @builder
     def replace_table(self, current_table, new_table):
@@ -834,18 +818,15 @@ class BitwiseAndCriterion(Criterion):
             term=self.term.get_sql(**kwargs), value=self.value,
         )
 
-    def fields(self):
-        return self.term.fields() if self.term.fields else []
-
 
 class NullCriterion(Criterion):
     def __init__(self, term, alias=None):
         super(NullCriterion, self).__init__(alias)
         self.term = term
 
-    @property
-    def tables_(self):
-        return self.term.tables_
+    def nodes_(self):
+        yield self
+        yield from self.term.nodes_()
 
     @builder
     def replace_table(self, current_table, new_table):
@@ -864,14 +845,8 @@ class NullCriterion(Criterion):
     def get_sql(self, **kwargs):
         return "{term} IS NULL".format(term=self.term.get_sql(**kwargs),)
 
-    def fields(self):
-        return self.term.fields() if self.term.fields else []
-
 
 class ComplexCriterion(BasicCriterion):
-    def fields(self):
-        return self.left.fields() + self.right.fields()
-
     def get_sql(self, subcriterion=False, **kwargs):
         sql = "{left} {comparator} {right}".format(
             comparator=self.comparator.value,
@@ -925,15 +900,16 @@ class ArithmeticExpression(Term):
         self.left = left
         self.right = right
 
+    def nodes_(self):
+        yield self
+        yield from self.left.nodes_()
+        yield from self.right.nodes_()
+
     @property
     def is_aggregate(self):
         # True if both left and right terms are True or None. None if both terms are None. Otherwise, False
         return resolve_is_aggregate([self.left.is_aggregate, self.right.is_aggregate])
 
-    @property
-    def tables_(self):
-        return self.left.tables_ | self.right.tables_
-
     @builder
     def replace_table(self, current_table, new_table):
         """
@@ -949,9 +925,6 @@ class ArithmeticExpression(Term):
         self.left = self.left.replace_table(current_table, new_table)
         self.right = self.right.replace_table(current_table, new_table)
 
-    def fields(self):
-        return self.left.fields() + self.right.fields()
-
     def get_sql(self, with_alias=False, **kwargs):
         is_mul = self.operator in self.mul_order
         is_left_add, is_right_add = [
@@ -959,8 +932,7 @@ class ArithmeticExpression(Term):
             for side in [self.left, self.right]
         ]
 
-        quote_char = kwargs.get("quote_char", None)
-        arithmatic_sql = "{left}{operator}{right}".format(
+        arithmetic_sql = "{left}{operator}{right}".format(
             operator=self.operator.value,
             left=("({})" if is_mul and is_left_add else "{}").format(
                 self.left.get_sql(**kwargs)
@@ -971,9 +943,9 @@ class ArithmeticExpression(Term):
         )
 
         if with_alias:
-            return format_alias_sql(arithmatic_sql, self.alias, **kwargs)
+            return format_alias_sql(arithmetic_sql, self.alias, **kwargs)
 
-        return arithmatic_sql
+        return arithmetic_sql
 
 
 class Case(Term):
@@ -982,6 +954,16 @@ class Case(Term):
         self._cases = []
         self._else = None
 
+    def nodes_(self):
+        yield self
+
+        for criterion, term in self._cases:
+            yield from criterion.nodes_()
+            yield from term.nodes_()
+
+        if self._else is not None:
+            yield from self._else.nodes_()
+
     @property
     def is_aggregate(self):
         # True if all cases are True or None. None all cases are None. Otherwise, False
@@ -1043,42 +1025,15 @@ class Case(Term):
 
         return case_sql
 
-    def fields(self):
-        fields = []
-
-        for criterion, term in self._cases:
-            fields += criterion.fields() + term.fields()
-
-        if self._else is not None:
-            fields += self._else.fields()
-
-        return fields
-
-    @property
-    def tables_(self):
-        tables = set()
-        if self._cases:
-            tables |= {
-                table
-                for case in self._cases
-                for part in case
-                for table in part.tables_
-                if hasattr(part, "tables_")
-            }
-
-        if self._else and hasattr(self._else, "tables_"):
-            tables |= {table for table in self._else.tables_}
-
-        return tables
-
 
 class Not(Criterion):
     def __init__(self, term, alias=None):
         super(Not, self).__init__(alias=alias)
         self.term = term
 
-    def fields(self):
-        return self.term.fields() if self.term.fields else []
+    def nodes_(self):
+        yield self
+        yield from self.term.nodes_()
 
     def get_sql(self, **kwargs):
         kwargs["subcriterion"] = True
@@ -1118,10 +1073,6 @@ class Not(Criterion):
         """
         self.term = self.term.replace_table(current_table, new_table)
 
-    @property
-    def tables_(self):
-        return self.term.tables_
-
 
 class CustomFunction:
     def __init__(self, name, params=None):
@@ -1157,17 +1108,10 @@ class Function(Criterion):
         self.args = [self.wrap_constant(param) for param in args]
         self.schema = kwargs.get("schema")
 
-    @property
-    def tables_(self):
-        return {table for param in self.args for table in param.tables_}
-
-    def fields(self):
-        return [
-            field
-            for param in self.args
-            if hasattr(param, "fields")
-            for field in param.fields()
-        ]
+    def nodes_(self):
+        yield self
+        for arg in self.args:
+            yield from arg.nodes_()
 
     @property
     def is_aggregate(self):
@@ -1375,7 +1319,7 @@ class IgnoreNullsAnalyticFunction(AnalyticFunction):
         return None
 
 
-class Interval:
+class Interval(Node):
     templates = {
         # MySQL requires no single quotes around the expr and unit
         Dialects.MYSQL: "INTERVAL {expr} {unit}",
@@ -1430,13 +1374,6 @@ class Interval:
     def __str__(self):
         return self.get_sql()
 
-    @property
-    def tables_(self):
-        return {}
-
-    def fields(self):
-        return []
-
     def get_sql(self, **kwargs):
         dialect = self.dialect or kwargs.get("dialect")
 
@@ -1504,6 +1441,3 @@ class PseudoColumn(Term):
 
     def get_sql(self, **kwargs):
         return self.name
-
-    def fields(self):
-        return []


Test patch:
diff --git a/pypika/tests/test_internals.py b/pypika/tests/test_internals.py
index 6d20361..921e381 100644
--- a/pypika/tests/test_internals.py
+++ b/pypika/tests/test_internals.py
@@ -26,7 +26,7 @@ class TablesTests(unittest.TestCase):
     def test__star_tables(self):
         star = Star()
 
-        self.assertEqual(star.tables_, {})
+        self.assertEqual(star.tables_, set())
 
     def test__table_star_tables(self):
         table = Table("a")