Define type hints for remaining column operators

mjpieters · zzzeek · commit ceec98e1d4d2 · 2023-04-18T11:38:13.000-04:00
Added typing information for recently added operators :meth:`.ColumnOperators.icontains`, :meth:`.ColumnOperators.istartswith`, :meth:`.ColumnOperators.iendswith`, and bitwise operators :meth:`.ColumnOperators.bitwise_and`, :meth:`.ColumnOperators.bitwise_or`, :meth:`.ColumnOperators.bitwise_xor`, :meth:`.ColumnOperators.bitwise_not`, :meth:`.ColumnOperators.bitwise_lshift` :meth:`.ColumnOperators.bitwise_rshift`. Pull request courtesy Martijn Pieters. Fixes: sqlalchemy#9650 Closes: sqlalchemy#9652 Pull-request: sqlalchemy#9652 Pull-request-sha: 005c568 Change-Id: I2fa06eb42ce668df9d9c760d233906f87484dd12
diff --git a/doc/build/changelog/unreleased_20/9650.rst b/doc/build/changelog/unreleased_20/9650.rst
@@ -0,0 +1,13 @@
+.. change::
+    :tags: bug, typing
+    :tickets: 9650
+
+    Added typing information for recently added operators
+    :meth:`.ColumnOperators.icontains`, :meth:`.ColumnOperators.istartswith`,
+    :meth:`.ColumnOperators.iendswith`, and bitwise operators
+    :meth:`.ColumnOperators.bitwise_and`, :meth:`.ColumnOperators.bitwise_or`,
+    :meth:`.ColumnOperators.bitwise_xor`, :meth:`.ColumnOperators.bitwise_not`,
+    :meth:`.ColumnOperators.bitwise_lshift`
+    :meth:`.ColumnOperators.bitwise_rshift`. Pull request courtesy Martijn
+    Pieters.
+
diff --git a/lib/sqlalchemy/sql/elements.py b/lib/sqlalchemy/sql/elements.py
@@ -833,6 +833,12 @@ def __contains__(self, other: Any) -> ColumnElement[bool]:
         def __getitem__(self, index: Any) -> ColumnElement[Any]:
             ...
 
+        def __lshift__(self, other: Any) -> ColumnElement[Any]:
+            ...
+
+        def __rshift__(self, other: Any) -> ColumnElement[Any]:
+            ...
+
         @overload
         def concat(self: _SQO[str], other: Any) -> ColumnElement[str]:
             ...
@@ -854,6 +860,24 @@ def ilike(
         ) -> BinaryExpression[bool]:
             ...
 
+        def bitwise_xor(self, other: Any) -> BinaryExpression[Any]:
+            ...
+
+        def bitwise_or(self, other: Any) -> BinaryExpression[Any]:
+            ...
+
+        def bitwise_and(self, other: Any) -> BinaryExpression[Any]:
+            ...
+
+        def bitwise_not(self) -> UnaryExpression[_T]:
+            ...
+
+        def bitwise_lshift(self, other: Any) -> BinaryExpression[Any]:
+            ...
+
+        def bitwise_rshift(self, other: Any) -> BinaryExpression[Any]:
+            ...
+
         def in_(
             self,
             other: Union[
@@ -915,6 +939,14 @@ def startswith(
         ) -> ColumnElement[bool]:
             ...
 
+        def istartswith(
+            self,
+            other: Any,
+            escape: Optional[str] = None,
+            autoescape: bool = False,
+        ) -> ColumnElement[bool]:
+            ...
+
         def endswith(
             self,
             other: Any,
@@ -923,9 +955,20 @@ def endswith(
         ) -> ColumnElement[bool]:
             ...
 
+        def iendswith(
+            self,
+            other: Any,
+            escape: Optional[str] = None,
+            autoescape: bool = False,
+        ) -> ColumnElement[bool]:
+            ...
+
         def contains(self, other: Any, **kw: Any) -> ColumnElement[bool]:
             ...
 
+        def icontains(self, other: Any, **kw: Any) -> ColumnElement[bool]:
+            ...
+
         def match(self, other: Any, **kwargs: Any) -> ColumnElement[bool]:
             ...
 
diff --git a/test/ext/mypy/plain_files/sql_operations.py b/test/ext/mypy/plain_files/sql_operations.py
@@ -67,23 +67,54 @@
 
 expr9 = c1.bool_op("@@")(func.to_tsquery("some & query"))
 
-# add tests for #9148
-and_(c1.is_(q))
-and_(c1.is_not(q))
-and_(c1.isnot(q))
-and_(c1.not_in(["x"]))
-and_(c1.notin_(["x"]))
-and_(c1.not_like("x"))
-and_(c1.notlike("x"))
-and_(c1.not_ilike("x"))
-and_(c1.notilike("x"))
-
-# issue #9451
-s1 = c1.cast(Integer)
-s2 = c1.cast(Float)
-s3 = c1.op("foobar")("operand").cast(DateTime)
-s4 = cast(c1, Float)
-s5 = cast(c1.op("foobar")("operand"), DateTime)
+
+def test_issue_9418() -> None:
+    and_(c1.is_(q))
+    and_(c1.is_not(q))
+    and_(c1.isnot(q))
+    and_(c1.not_in(["x"]))
+    and_(c1.notin_(["x"]))
+    and_(c1.not_like("x"))
+    and_(c1.notlike("x"))
+    and_(c1.not_ilike("x"))
+    and_(c1.notilike("x"))
+
+
+def test_issue_9451() -> None:
+    # issue #9451
+    c1.cast(Integer)
+    c1.cast(Float)
+    c1.op("foobar")("operand").cast(DateTime)
+    cast(c1, Float)
+    cast(c1.op("foobar")("operand"), DateTime)
+
+
+def test_issue_9650_char() -> None:
+    and_(c1.contains("x"))
+    and_(c1.startswith("x"))
+    and_(c1.endswith("x"))
+    and_(c1.icontains("x"))
+    and_(c1.istartswith("x"))
+    and_(c1.iendswith("x"))
+
+
+def test_issue_9650_bitwise() -> None:
+    # EXPECTED_TYPE: BinaryExpression[Any]
+    reveal_type(c2.bitwise_and(5))
+    # EXPECTED_TYPE: BinaryExpression[Any]
+    reveal_type(c2.bitwise_or(5))
+    # EXPECTED_TYPE: BinaryExpression[Any]
+    reveal_type(c2.bitwise_xor(5))
+    # EXPECTED_TYPE: UnaryExpression[int]
+    reveal_type(c2.bitwise_not())
+    # EXPECTED_TYPE: BinaryExpression[Any]
+    reveal_type(c2.bitwise_lshift(5))
+    # EXPECTED_TYPE: BinaryExpression[Any]
+    reveal_type(c2.bitwise_rshift(5))
+    # EXPECTED_TYPE: ColumnElement[Any]
+    reveal_type(c2 << 5)
+    # EXPECTED_TYPE: ColumnElement[Any]
+    reveal_type(c2 >> 5)
 
 
 if typing.TYPE_CHECKING:
