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 python/mypy#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 microsoft/pyright#2758 (comment) . 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
sqlalchemy · Dec 30, 2021 · e913ec8 · e913ec8
1 parent 54875c2
commit e913ec8
Show file tree

Hide file tree

Showing 21 changed files with 484 additions and 158 deletions.
diff --git a/MANIFEST.in b/MANIFEST.in
@@ -7,7 +7,7 @@ recursive-include test *.py *.dat *.testpatch
 
 # include the pyx and pxd extensions, which otherwise
 # don't come in if --with-cextensions isn't specified.
-recursive-include lib *.pyx *.pxd *.txt
+recursive-include lib *.pyx *.pxd *.txt *.typed
 
 include README* AUTHORS LICENSE CHANGES* tox.ini
 prune doc/build/output
diff --git a/lib/sqlalchemy/dialects/mysql/dml.py b/lib/sqlalchemy/dialects/mysql/dml.py
@@ -1,3 +1,11 @@
+# Copyright (C) 2005-2021 the SQLAlchemy authors and contributors
+# <see AUTHORS file>
+#
+# This module is part of SQLAlchemy and is released under
+# the MIT License: https://www.opensource.org/licenses/mit-license.php
+
+import typing
+
 from ... import exc
 from ... import util
 from ...sql.base import _exclusive_against
@@ -12,6 +20,9 @@
 __all__ = ("Insert", "insert")
 
 
+SelfInsert = typing.TypeVar("SelfInsert", bound="Insert")
+
+
 class Insert(StandardInsert):
     """MySQL-specific implementation of INSERT.
 
@@ -70,7 +81,7 @@ def inserted_alias(self):
             "has an ON DUPLICATE KEY clause present"
         },
     )
-    def on_duplicate_key_update(self, *args, **kw):
+    def on_duplicate_key_update(self: SelfInsert, *args, **kw) -> SelfInsert:
         r"""
         Specifies the ON DUPLICATE KEY UPDATE clause.
 
@@ -131,6 +142,7 @@ def on_duplicate_key_update(self, *args, **kw):
 
         inserted_alias = getattr(self, "inserted_alias", None)
         self._post_values_clause = OnDuplicateClause(inserted_alias, values)
+        return self
 
 
 insert = public_factory(

diff --git a/lib/sqlalchemy/dialects/mysql/expression.py b/lib/sqlalchemy/dialects/mysql/expression.py
@@ -1,3 +1,11 @@
+# Copyright (C) 2005-2021 the SQLAlchemy authors and contributors
+# <see AUTHORS file>
+#
+# This module is part of SQLAlchemy and is released under
+# the MIT License: https://www.opensource.org/licenses/mit-license.php
+
+import typing
+
 from ... import exc
 from ... import util
 from ...sql import coercions
@@ -8,6 +16,9 @@
 from ...sql.base import Generative
 
 
+Selfmatch = typing.TypeVar("Selfmatch", bound="match")
+
+
 class match(Generative, elements.BinaryExpression):
     """Produce a ``MATCH (X, Y) AGAINST ('TEXT')`` clause.
 
@@ -99,17 +110,18 @@ def __init__(self, *cols, **kw):
         )
 
     @_generative
-    def in_boolean_mode(self):
+    def in_boolean_mode(self: Selfmatch) -> Selfmatch:
         """Apply the "IN BOOLEAN MODE" modifier to the MATCH expression.
 
         :return: a new :class:`_mysql.match` instance with modifications
          applied.
         """
 
         self.modifiers = self.modifiers.union({"mysql_boolean_mode": True})
+        return self
 
     @_generative
-    def in_natural_language_mode(self):
+    def in_natural_language_mode(self: Selfmatch) -> Selfmatch:
         """Apply the "IN NATURAL LANGUAGE MODE" modifier to the MATCH
         expression.
 
@@ -118,13 +130,15 @@ def in_natural_language_mode(self):
         """
 
         self.modifiers = self.modifiers.union({"mysql_natural_language": True})
+        return self
 
     @_generative
-    def with_query_expansion(self):
+    def with_query_expansion(self: Selfmatch) -> Selfmatch:
         """Apply the "WITH QUERY EXPANSION" modifier to the MATCH expression.
 
         :return: a new :class:`_mysql.match` instance with modifications
          applied.
         """
 
         self.modifiers = self.modifiers.union({"mysql_query_expansion": True})
+        return self
diff --git a/lib/sqlalchemy/dialects/postgresql/dml.py b/lib/sqlalchemy/dialects/postgresql/dml.py
@@ -4,6 +4,7 @@
 #
 # This module is part of SQLAlchemy and is released under
 # the MIT License: https://www.opensource.org/licenses/mit-license.php
+import typing
 
 from . import ext
 from ... import util
@@ -21,6 +22,8 @@
 
 __all__ = ("Insert", "insert")
 
+SelfInsert = typing.TypeVar("SelfInsert", bound="Insert")
+
 
 class Insert(StandardInsert):
     """PostgreSQL-specific implementation of INSERT.
@@ -75,13 +78,13 @@ def excluded(self):
     @_generative
     @_on_conflict_exclusive
     def on_conflict_do_update(
-        self,
+        self: SelfInsert,
         constraint=None,
         index_elements=None,
         index_where=None,
         set_=None,
         where=None,
-    ):
+    ) -> SelfInsert:
         r"""
         Specifies a DO UPDATE SET action for ON CONFLICT clause.
 
@@ -138,12 +141,16 @@ def on_conflict_do_update(
         self._post_values_clause = OnConflictDoUpdate(
             constraint, index_elements, index_where, set_, where
         )
+        return self
 
     @_generative
     @_on_conflict_exclusive
     def on_conflict_do_nothing(
-        self, constraint=None, index_elements=None, index_where=None
-    ):
+        self: SelfInsert,
+        constraint=None,
+        index_elements=None,
+        index_where=None,
+    ) -> SelfInsert:
         """
         Specifies a DO NOTHING action for ON CONFLICT clause.
 
@@ -173,6 +180,7 @@ def on_conflict_do_nothing(
         self._post_values_clause = OnConflictDoNothing(
             constraint, index_elements, index_where
         )
+        return self
 
 
 insert = public_factory(

diff --git a/lib/sqlalchemy/dialects/sqlite/dml.py b/lib/sqlalchemy/dialects/sqlite/dml.py
@@ -4,6 +4,8 @@
 # This module is part of SQLAlchemy and is released under
 # the MIT License: https://www.opensource.org/licenses/mit-license.php
 
+import typing
+
 from ... import util
 from ...sql import coercions
 from ...sql import roles
@@ -18,6 +20,8 @@
 
 __all__ = ("Insert", "insert")
 
+SelfInsert = typing.TypeVar("SelfInsert", bound="Insert")
+
 
 class Insert(StandardInsert):
     """SQLite-specific implementation of INSERT.
@@ -71,12 +75,12 @@ def excluded(self):
     @_generative
     @_on_conflict_exclusive
     def on_conflict_do_update(
-        self,
+        self: SelfInsert,
         index_elements=None,
         index_where=None,
         set_=None,
         where=None,
-    ):
+    ) -> SelfInsert:
         r"""
         Specifies a DO UPDATE SET action for ON CONFLICT clause.
 
@@ -120,10 +124,13 @@ def on_conflict_do_update(
         self._post_values_clause = OnConflictDoUpdate(
             index_elements, index_where, set_, where
         )
+        return self
 
     @_generative
     @_on_conflict_exclusive
-    def on_conflict_do_nothing(self, index_elements=None, index_where=None):
+    def on_conflict_do_nothing(
+        self: SelfInsert, index_elements=None, index_where=None
+    ) -> SelfInsert:
         """
         Specifies a DO NOTHING action for ON CONFLICT clause.
 
@@ -141,6 +148,7 @@ def on_conflict_do_nothing(self, index_elements=None, index_where=None):
         self._post_values_clause = OnConflictDoNothing(
             index_elements, index_where
         )
+        return self
 
 
 insert = public_factory(

diff --git a/lib/sqlalchemy/engine/cursor.py b/lib/sqlalchemy/engine/cursor.py
@@ -1702,6 +1702,7 @@ def close(self):
     def yield_per(self, num):
         self._yield_per = num
         self.cursor_strategy.yield_per(self, self.cursor, num)
+        return self
 
 
 ResultProxy = CursorResult

diff --git a/lib/sqlalchemy/engine/result.py b/lib/sqlalchemy/engine/result.py
@@ -7,11 +7,11 @@
 
 """Define generic result set constructs."""
 
-
 import collections.abc as collections_abc
 import functools
 import itertools
 import operator
+import typing
 
 from .row import Row
 from .. import exc
@@ -257,6 +257,10 @@ def result_tuple(fields, extra=None):
 # filter is applied to rows.
 _NO_ROW = util.symbol("NO_ROW")
 
+SelfResultInternal = typing.TypeVar(
+    "SelfResultInternal", bound="ResultInternal"
+)
+
 
 class ResultInternal(InPlaceGenerative):
     _real_result = None
@@ -614,7 +618,9 @@ def _next_impl(self):
             return row
 
     @_generative
-    def _column_slices(self, indexes):
+    def _column_slices(
+        self: SelfResultInternal, indexes
+    ) -> SelfResultInternal:
         real_result = self._real_result if self._real_result else self
 
         if real_result._source_supports_scalars and len(indexes) == 1:
@@ -623,6 +629,8 @@ def _column_slices(self, indexes):
             self._generate_rows = True
             self._metadata = self._metadata._reduce(indexes)
 
+        return self
+
     @HasMemoized.memoized_attribute
     def _unique_strategy(self):
         uniques, strategy = self._unique_filter_state
@@ -668,6 +676,9 @@ def keys(self):
         return self._metadata.keys
 
 
+SelfResult = typing.TypeVar("SelfResult", bound="Result")
+
+
 class Result(_WithKeys, ResultInternal):
     """Represent a set of database results.
 
@@ -732,7 +743,7 @@ def close(self):
         self._soft_close(hard=True)
 
     @_generative
-    def yield_per(self, num):
+    def yield_per(self: SelfResult, num) -> SelfResult:
         """Configure the row-fetching strategy to fetch num rows at a time.
 
         This impacts the underlying behavior of the result when iterating over
@@ -766,9 +777,10 @@ def yield_per(self, num):
 
         """
         self._yield_per = num
+        return self
 
     @_generative
-    def unique(self, strategy=None):
+    def unique(self: SelfResult, strategy=None) -> SelfResult:
         """Apply unique filtering to the objects returned by this
         :class:`_engine.Result`.
 
@@ -806,8 +818,11 @@ def unique(self, strategy=None):
 
         """
         self._unique_filter_state = (set(), strategy)
+        return self
 
-    def columns(self, *col_expressions):
+    def columns(
+        self: SelfResultInternal, *col_expressions
+    ) -> SelfResultInternal:
         r"""Establish the columns that should be returned in each row.
 
         This method may be used to limit the columns returned as well
@@ -845,7 +860,7 @@ def columns(self, *col_expressions):
         """
         return self._column_slices(col_expressions)
 
-    def scalars(self, index=0):
+    def scalars(self, index=0) -> "ScalarResult":
         """Return a :class:`_result.ScalarResult` filtering object which
         will return single elements rather than :class:`_row.Row` objects.
 
@@ -892,7 +907,7 @@ def _tuple_getter(self, keys):
             )
         return self._metadata._row_as_tuple_getter(keys)
 
-    def mappings(self):
+    def mappings(self) -> "MappingResult":
         """Apply a mappings filter to returned rows, returning an instance of
         :class:`_result.MappingResult`.
 
@@ -1653,6 +1668,11 @@ def null_result():
     return IteratorResult(SimpleResultMetaData([]), iter([]))
 
 
+SelfChunkedIteratorResult = typing.TypeVar(
+    "SelfChunkedIteratorResult", bound="ChunkedIteratorResult"
+)
+
+
 class ChunkedIteratorResult(IteratorResult):
     """An :class:`.IteratorResult` that works from an iterator-producing callable.
 
@@ -1684,7 +1704,9 @@ def __init__(
         self.dynamic_yield_per = dynamic_yield_per
 
     @_generative
-    def yield_per(self, num):
+    def yield_per(
+        self: SelfChunkedIteratorResult, num
+    ) -> SelfChunkedIteratorResult:
         # TODO: this throws away the iterator which may be holding
         # onto a chunk.   the yield_per cannot be changed once any
         # rows have been fetched.   either find a way to enforce this,
@@ -1693,6 +1715,7 @@ def yield_per(self, num):
 
         self._yield_per = num
         self.iterator = itertools.chain.from_iterable(self.chunks(num))
+        return self
 
     def _soft_close(self, **kw):
         super(ChunkedIteratorResult, self)._soft_close(**kw)