Properly type _generative, decorator, public_factory

Good new is that pylance likes it and copies over the
singature and everything.
Bad news is that mypy does not support this yet https://github.com/python/mypy/issues/8645
Other minor bad news is that non_generative is not typed. I've tried using a protocol
like the one in the comment but the signature is not ported over by pylance, so it's
probably best to just live without it to have the correct signature.

notes from mike:  these three decorators are at the core of getting
the library to be typed, more good news is that pylance will
do all the things we like re: public_factory, see
https://github.com/microsoft/pyright/issues/2758#issuecomment-1002788656
.

For @_generative, we will likely move to using pep 673 once mypy
supports it which may be soon.  but overall having the explicit
"return self" in the methods, while a little inconvenient, makes
the typing more straightforward and locally present in the files
rather than being decided at a distance.   having "return self"
present, or not, both have problems, so maybe we will be able
to change it again if things change as far as decorator support.
As it is, I feel like we are barely squeaking by with our decorators,
the typing is already pretty out there.

Change-Id: Ic77e13fc861def76a5925331df85c0aa48d77807
References: #6810

1 files changed, 110 insertions, 32 deletions

diff --git a/lib/sqlalchemy/sql/selectable.py b/lib/sqlalchemy/sql/selectable.py
index 802576b89..8b35dc6ac 100644
--- a/lib/sqlalchemy/sql/selectable.py
+++ b/lib/sqlalchemy/sql/selectable.py
@@ -14,6 +14,9 @@ SQL tables and derived rowsets.
 import collections
 import itertools
 from operator import attrgetter
+import typing
+from typing import Type
+from typing import Union
 
 from . import coercions
 from . import operators
@@ -209,6 +212,9 @@ class Selectable(ReturnsRows):
         )
 
 
+SelfHasPrefixes = typing.TypeVar("SelfHasPrefixes", bound="HasPrefixes")
+
+
 class HasPrefixes:
     _prefixes = ()
 
@@ -222,7 +228,7 @@ class HasPrefixes:
         ":meth:`_expression.HasPrefixes.prefix_with`",
         ":paramref:`.HasPrefixes.prefix_with.*expr`",
     )
-    def prefix_with(self, *expr, **kw):
+    def prefix_with(self: SelfHasPrefixes, *expr, **kw) -> SelfHasPrefixes:
         r"""Add one or more expressions following the statement keyword, i.e.
         SELECT, INSERT, UPDATE, or DELETE. Generative.
 
@@ -255,6 +261,7 @@ class HasPrefixes:
                 "Unsupported argument(s): %s" % ",".join(kw)
             )
         self._setup_prefixes(expr, dialect)
+        return self
 
     def _setup_prefixes(self, prefixes, dialect=None):
         self._prefixes = self._prefixes + tuple(
@@ -265,6 +272,9 @@ class HasPrefixes:
         )
 
 
+SelfHasSuffixes = typing.TypeVar("SelfHasSuffixes", bound="HasSuffixes")
+
+
 class HasSuffixes:
     _suffixes = ()
 
@@ -278,7 +288,7 @@ class HasSuffixes:
         ":meth:`_expression.HasSuffixes.suffix_with`",
         ":paramref:`.HasSuffixes.suffix_with.*expr`",
     )
-    def suffix_with(self, *expr, **kw):
+    def suffix_with(self: SelfHasSuffixes, *expr, **kw) -> SelfHasSuffixes:
         r"""Add one or more expressions following the statement as a whole.
 
         This is used to support backend-specific suffix keywords on
@@ -306,6 +316,7 @@ class HasSuffixes:
                 "Unsupported argument(s): %s" % ",".join(kw)
             )
         self._setup_suffixes(expr, dialect)
+        return self
 
     def _setup_suffixes(self, suffixes, dialect=None):
         self._suffixes = self._suffixes + tuple(
@@ -316,6 +327,9 @@ class HasSuffixes:
         )
 
 
+SelfHasHints = typing.TypeVar("SelfHasHints", bound="HasHints")
+
+
 class HasHints:
     _hints = util.immutabledict()
     _statement_hints = ()
@@ -352,7 +366,9 @@ class HasHints:
         return self.with_hint(None, text, dialect_name)
 
     @_generative
-    def with_hint(self, selectable, text, dialect_name="*"):
+    def with_hint(
+        self: SelfHasHints, selectable, text, dialect_name="*"
+    ) -> SelfHasHints:
         r"""Add an indexing or other executional context hint for the given
         selectable to this :class:`_expression.Select` or other selectable
         object.
@@ -398,6 +414,7 @@ class HasHints:
                     ): text
                 }
             )
+        return self
 
 
 class FromClause(roles.AnonymizedFromClauseRole, Selectable):
@@ -2082,6 +2099,9 @@ class CTE(
         return self._restates if self._restates is not None else self
 
 
+SelfHasCTE = typing.TypeVar("SelfHasCTE", bound="HasCTE")
+
+
 class HasCTE(roles.HasCTERole):
     """Mixin that declares a class to include CTE support.
 
@@ -2096,7 +2116,7 @@ class HasCTE(roles.HasCTERole):
     _independent_ctes = ()
 
     @_generative
-    def add_cte(self, cte):
+    def add_cte(self: SelfHasCTE, cte) -> SelfHasCTE:
         """Add a :class:`_sql.CTE` to this statement object that will be
         independently rendered even if not referenced in the statement
         otherwise.
@@ -2161,6 +2181,7 @@ class HasCTE(roles.HasCTERole):
         """
         cte = coercions.expect(roles.IsCTERole, cte)
         self._independent_ctes += (cte,)
+        return self
 
     def cte(self, name=None, recursive=False, nesting=False):
         r"""Return a new :class:`_expression.CTE`,
@@ -2759,6 +2780,9 @@ class ForUpdateArg(ClauseElement):
             self.of = None
 
 
+SelfValues = typing.TypeVar("SelfValues", bound="Values")
+
+
 class Values(Generative, FromClause):
     """Represent a ``VALUES`` construct that can be used as a FROM element
     in a statement.
@@ -2829,7 +2853,7 @@ class Values(Generative, FromClause):
         return [col.type for col in self._column_args]
 
     @_generative
-    def alias(self, name, **kw):
+    def alias(self: SelfValues, name, **kw) -> SelfValues:
         """Return a new :class:`_expression.Values`
         construct that is a copy of this
         one with the given name.
@@ -2846,9 +2870,10 @@ class Values(Generative, FromClause):
         """
         self.name = name
         self.named_with_column = self.name is not None
+        return self
 
     @_generative
-    def lateral(self, name=None):
+    def lateral(self: SelfValues, name=None) -> SelfValues:
         """Return a new :class:`_expression.Values` with the lateral flag set,
         so that
         it renders as LATERAL.
@@ -2861,9 +2886,10 @@ class Values(Generative, FromClause):
         self._is_lateral = True
         if name is not None:
             self.name = name
+        return self
 
     @_generative
-    def data(self, values):
+    def data(self: SelfValues, values) -> SelfValues:
         """Return a new :class:`_expression.Values` construct,
         adding the given data
         to the data list.
@@ -2879,6 +2905,7 @@ class Values(Generative, FromClause):
         """
 
         self._data += (values,)
+        return self
 
     def _populate_column_collection(self):
         for c in self._column_args:
@@ -3312,6 +3339,11 @@ class DeprecatedSelectBaseGenerations:
         self.group_by.non_generative(self, *clauses)
 
 
+SelfGenerativeSelect = typing.TypeVar(
+    "SelfGenerativeSelect", bound="GenerativeSelect"
+)
+
+
 class GenerativeSelect(DeprecatedSelectBaseGenerations, SelectBase):
     """Base class for SELECT statements where additional elements can be
     added.
@@ -3371,13 +3403,13 @@ class GenerativeSelect(DeprecatedSelectBaseGenerations, SelectBase):
 
     @_generative
     def with_for_update(
-        self,
+        self: SelfGenerativeSelect,
         nowait=False,
         read=False,
         of=None,
         skip_locked=False,
         key_share=False,
-    ):
+    ) -> SelfGenerativeSelect:
         """Specify a ``FOR UPDATE`` clause for this
         :class:`_expression.GenerativeSelect`.
 
@@ -3430,6 +3462,7 @@ class GenerativeSelect(DeprecatedSelectBaseGenerations, SelectBase):
             skip_locked=skip_locked,
             key_share=key_share,
         )
+        return self
 
     def get_label_style(self):
         """
@@ -3573,7 +3606,7 @@ class GenerativeSelect(DeprecatedSelectBaseGenerations, SelectBase):
         )
 
     @_generative
-    def limit(self, limit):
+    def limit(self: SelfGenerativeSelect, limit) -> SelfGenerativeSelect:
         """Return a new selectable with the given LIMIT criterion
         applied.
 
@@ -3603,9 +3636,12 @@ class GenerativeSelect(DeprecatedSelectBaseGenerations, SelectBase):
 
         self._fetch_clause = self._fetch_clause_options = None
         self._limit_clause = self._offset_or_limit_clause(limit)
+        return self
 
     @_generative
-    def fetch(self, count, with_ties=False, percent=False):
+    def fetch(
+        self: SelfGenerativeSelect, count, with_ties=False, percent=False
+    ) -> SelfGenerativeSelect:
         """Return a new selectable with the given FETCH FIRST criterion
         applied.
 
@@ -3653,9 +3689,10 @@ class GenerativeSelect(DeprecatedSelectBaseGenerations, SelectBase):
                 "with_ties": with_ties,
                 "percent": percent,
             }
+        return self
 
     @_generative
-    def offset(self, offset):
+    def offset(self: SelfGenerativeSelect, offset) -> SelfGenerativeSelect:
         """Return a new selectable with the given OFFSET criterion
         applied.
 
@@ -3681,10 +3718,11 @@ class GenerativeSelect(DeprecatedSelectBaseGenerations, SelectBase):
         """
 
         self._offset_clause = self._offset_or_limit_clause(offset)
+        return self
 
     @_generative
     @util.preload_module("sqlalchemy.sql.util")
-    def slice(self, start, stop):
+    def slice(self: SelfGenerativeSelect, start, stop) -> SelfGenerativeSelect:
         """Apply LIMIT / OFFSET to this statement based on a slice.
 
         The start and stop indices behave like the argument to Python's
@@ -3728,9 +3766,10 @@ class GenerativeSelect(DeprecatedSelectBaseGenerations, SelectBase):
         self._limit_clause, self._offset_clause = sql_util._make_slice(
             self._limit_clause, self._offset_clause, start, stop
         )
+        return self
 
     @_generative
-    def order_by(self, *clauses):
+    def order_by(self: SelfGenerativeSelect, *clauses) -> SelfGenerativeSelect:
         r"""Return a new selectable with the given list of ORDER BY
         criteria applied.
 
@@ -3764,9 +3803,10 @@ class GenerativeSelect(DeprecatedSelectBaseGenerations, SelectBase):
                 coercions.expect(roles.OrderByRole, clause)
                 for clause in clauses
             )
+        return self
 
     @_generative
-    def group_by(self, *clauses):
+    def group_by(self: SelfGenerativeSelect, *clauses) -> SelfGenerativeSelect:
         r"""Return a new selectable with the given list of GROUP BY
         criterion applied.
 
@@ -3797,6 +3837,7 @@ class GenerativeSelect(DeprecatedSelectBaseGenerations, SelectBase):
                 coercions.expect(roles.GroupByRole, clause)
                 for clause in clauses
             )
+        return self
 
 
 @CompileState.plugin_for("default", "compound_select")
@@ -4658,6 +4699,10 @@ class _MemoizedSelectEntities(
             ) = select_stmt._with_options = ()
 
 
+# TODO: use pep-673 when feasible
+SelfSelect = typing.TypeVar("SelfSelect", bound="Select")
+
+
 class Select(
     HasPrefixes,
     HasSuffixes,
@@ -4737,7 +4782,9 @@ class Select(
     ]
 
     @classmethod
-    def _create(cls, *entities) -> "Select":
+    def _create(
+        cls, *entities: Union[roles.ColumnsClauseRole, Type]
+    ) -> "Select":
         r"""Construct a new :class:`_expression.Select`.
 
 
@@ -4788,7 +4835,7 @@ class Select(
         return self
 
     @classmethod
-    def _create_raw_select(cls, **kw):
+    def _create_raw_select(cls, **kw) -> "Select":
         """Create a :class:`.Select` using raw ``__new__`` with no coercions.
 
         Used internally to build up :class:`.Select` constructs with
@@ -4873,7 +4920,9 @@ class Select(
         return meth(self, statement)
 
     @_generative
-    def join(self, target, onclause=None, isouter=False, full=False):
+    def join(
+        self: SelfSelect, target, onclause=None, isouter=False, full=False
+    ) -> SelfSelect:
         r"""Create a SQL JOIN against this :class:`_expression.Select`
         object's criterion
         and apply generatively, returning the newly resulting
@@ -4939,6 +4988,7 @@ class Select(
         self._setup_joins += (
             (target, onclause, None, {"isouter": isouter, "full": full}),
         )
+        return self
 
     def outerjoin_from(self, from_, target, onclause=None, full=False):
         r"""Create a SQL LEFT OUTER JOIN against this :class:`_expression.Select`
@@ -4955,8 +5005,13 @@ class Select(
 
     @_generative
     def join_from(
-        self, from_, target, onclause=None, isouter=False, full=False
-    ):
+        self: SelfSelect,
+        from_,
+        target,
+        onclause=None,
+        isouter=False,
+        full=False,
+    ) -> SelfSelect:
         r"""Create a SQL JOIN against this :class:`_expression.Select`
         object's criterion
         and apply generatively, returning the newly resulting
@@ -5014,6 +5069,7 @@ class Select(
         self._setup_joins += (
             (target, onclause, from_, {"isouter": isouter, "full": full}),
         )
+        return self
 
     def outerjoin(self, target, onclause=None, full=False):
         """Create a left outer join.
@@ -5211,7 +5267,7 @@ class Select(
         )
 
     @_generative
-    def add_columns(self, *columns):
+    def add_columns(self: SelfSelect, *columns) -> SelfSelect:
         """Return a new :func:`_expression.select` construct with
         the given column expressions added to its columns clause.
 
@@ -5233,6 +5289,7 @@ class Select(
             )
             for column in columns
         ]
+        return self
 
     def _set_entities(self, entities):
         self._raw_columns = [
@@ -5297,7 +5354,7 @@ class Select(
         )
 
     @_generative
-    def with_only_columns(self, *columns, **kw):
+    def with_only_columns(self: SelfSelect, *columns, **kw) -> SelfSelect:
         r"""Return a new :func:`_expression.select` construct with its columns
         clause replaced with the given columns.
 
@@ -5372,6 +5429,7 @@ class Select(
                 "columns", "Select.with_only_columns", columns
             )
         ]
+        return self
 
     @property
     def whereclause(self):
@@ -5393,7 +5451,7 @@ class Select(
     _whereclause = whereclause
 
     @_generative
-    def where(self, *whereclause):
+    def where(self: SelfSelect, *whereclause) -> SelfSelect:
         """Return a new :func:`_expression.select` construct with
         the given expression added to
         its WHERE clause, joined to the existing clause via AND, if any.
@@ -5405,9 +5463,10 @@ class Select(
         for criterion in whereclause:
             where_criteria = coercions.expect(roles.WhereHavingRole, criterion)
             self._where_criteria += (where_criteria,)
+        return self
 
     @_generative
-    def having(self, having):
+    def having(self: SelfSelect, having) -> SelfSelect:
         """Return a new :func:`_expression.select` construct with
         the given expression added to
         its HAVING clause, joined to the existing clause via AND, if any.
@@ -5416,9 +5475,10 @@ class Select(
         self._having_criteria += (
             coercions.expect(roles.WhereHavingRole, having),
         )
+        return self
 
     @_generative
-    def distinct(self, *expr):
+    def distinct(self: SelfSelect, *expr) -> SelfSelect:
         r"""Return a new :func:`_expression.select` construct which
         will apply DISTINCT to its columns clause.
 
@@ -5437,9 +5497,10 @@ class Select(
             )
         else:
             self._distinct = True
+        return self
 
     @_generative
-    def select_from(self, *froms):
+    def select_from(self: SelfSelect, *froms) -> SelfSelect:
         r"""Return a new :func:`_expression.select` construct with the
         given FROM expression(s)
         merged into its list of FROM objects.
@@ -5480,9 +5541,10 @@ class Select(
             )
             for fromclause in froms
         )
+        return self
 
     @_generative
-    def correlate(self, *fromclauses):
+    def correlate(self: SelfSelect, *fromclauses) -> SelfSelect:
         r"""Return a new :class:`_expression.Select`
         which will correlate the given FROM
         clauses to that of an enclosing :class:`_expression.Select`.
@@ -5541,9 +5603,10 @@ class Select(
             self._correlate = self._correlate + tuple(
                 coercions.expect(roles.FromClauseRole, f) for f in fromclauses
             )
+        return self
 
     @_generative
-    def correlate_except(self, *fromclauses):
+    def correlate_except(self: SelfSelect, *fromclauses) -> SelfSelect:
         r"""Return a new :class:`_expression.Select`
         which will omit the given FROM
         clauses from the auto-correlation process.
@@ -5579,6 +5642,7 @@ class Select(
             self._correlate_except = (self._correlate_except or ()) + tuple(
                 coercions.expect(roles.FromClauseRole, f) for f in fromclauses
             )
+        return self
 
     @HasMemoized.memoized_attribute
     def selected_columns(self):
@@ -5959,6 +6023,9 @@ class Select(
         return CompoundSelect._create_intersect_all(self, *other, **kwargs)
 
 
+SelfScalarSelect = typing.TypeVar("SelfScalarSelect", bound="ScalarSelect")
+
+
 class ScalarSelect(roles.InElementRole, Generative, Grouping):
     """Represent a scalar subquery.
 
@@ -5998,18 +6065,19 @@ class ScalarSelect(roles.InElementRole, Generative, Grouping):
     c = columns
 
     @_generative
-    def where(self, crit):
+    def where(self: SelfScalarSelect, crit) -> SelfScalarSelect:
         """Apply a WHERE clause to the SELECT statement referred to
         by this :class:`_expression.ScalarSelect`.
 
         """
         self.element = self.element.where(crit)
+        return self
 
     def self_group(self, **kwargs):
         return self
 
     @_generative
-    def correlate(self, *fromclauses):
+    def correlate(self: SelfScalarSelect, *fromclauses) -> SelfScalarSelect:
         r"""Return a new :class:`_expression.ScalarSelect`
         which will correlate the given FROM
         clauses to that of an enclosing :class:`_expression.Select`.
@@ -6039,9 +6107,12 @@ class ScalarSelect(roles.InElementRole, Generative, Grouping):
 
         """
         self.element = self.element.correlate(*fromclauses)
+        return self
 
     @_generative
-    def correlate_except(self, *fromclauses):
+    def correlate_except(
+        self: SelfScalarSelect, *fromclauses
+    ) -> SelfScalarSelect:
         r"""Return a new :class:`_expression.ScalarSelect`
         which will omit the given FROM
         clauses from the auto-correlation process.
@@ -6073,6 +6144,7 @@ class ScalarSelect(roles.InElementRole, Generative, Grouping):
         """
 
         self.element = self.element.correlate_except(*fromclauses)
+        return self
 
 
 class Exists(UnaryExpression):
@@ -6228,6 +6300,9 @@ class Exists(UnaryExpression):
         return e
 
 
+SelfTextualSelect = typing.TypeVar("SelfTextualSelect", bound="TextualSelect")
+
+
 class TextualSelect(SelectBase):
     """Wrap a :class:`_expression.TextClause` construct within a
     :class:`_expression.SelectBase`
@@ -6315,8 +6390,11 @@ class TextualSelect(SelectBase):
         return self
 
     @_generative
-    def bindparams(self, *binds, **bind_as_values):
+    def bindparams(
+        self: SelfTextualSelect, *binds, **bind_as_values
+    ) -> SelfTextualSelect:
         self.element = self.element.bindparams(*binds, **bind_as_values)
+        return self
 
     def _generate_fromclause_column_proxies(self, fromclause):
         fromclause._columns._populate_separate_keys(