Fixed type annotations in _BaseLineQid and _BaseGridQid (quantumlib#6043

)
TarunSinghania · Mar 21, 2023 · 7335e71 · 7335e71
1 parent f636c5f
commit 7335e71
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/cirq-core/cirq/devices/grid_qubit.py b/cirq-core/cirq/devices/grid_qubit.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import functools
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Set, TypeVar, TYPE_CHECKING
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Set, TypeVar, TYPE_CHECKING, Union
 
 import abc
 
@@ -75,7 +75,7 @@ def _with_row_col(self: TSelf, row: int, col: int) -> TSelf:
     def __complex__(self) -> complex:
         return self.col + 1j * self.row
 
-    def __add__(self: TSelf, other: Tuple[int, int]) -> 'TSelf':
+    def __add__(self: TSelf, other: Union[Tuple[int, int], TSelf]) -> 'TSelf':
         if isinstance(other, _BaseGridQid):
             if self.dimension != other.dimension:
                 raise TypeError(
@@ -94,7 +94,7 @@ def __add__(self: TSelf, other: Tuple[int, int]) -> 'TSelf':
             )
         return self._with_row_col(row=self.row + other[0], col=self.col + other[1])
 
-    def __sub__(self: TSelf, other: Tuple[int, int]) -> 'TSelf':
+    def __sub__(self: TSelf, other: Union[Tuple[int, int], TSelf]) -> 'TSelf':
         if isinstance(other, _BaseGridQid):
             if self.dimension != other.dimension:
                 raise TypeError(

diff --git a/cirq-core/cirq/devices/line_qubit.py b/cirq-core/cirq/devices/line_qubit.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import functools
-from typing import Any, Dict, Iterable, List, Optional, Sequence, Set, TypeVar, TYPE_CHECKING
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Set, TypeVar, TYPE_CHECKING, Union
 
 import abc
 
@@ -69,7 +69,7 @@ def neighbors(self, qids: Optional[Iterable[ops.Qid]] = None) -> Set['_BaseLineQ
     def _with_x(self: TSelf, x: int) -> TSelf:
         """Returns a qubit with the same type but a different value of `x`."""
 
-    def __add__(self: TSelf, other: int) -> TSelf:
+    def __add__(self: TSelf, other: Union[int, TSelf]) -> TSelf:
         if isinstance(other, _BaseLineQid):
             if self.dimension != other.dimension:
                 raise TypeError(
@@ -81,7 +81,7 @@ def __add__(self: TSelf, other: int) -> TSelf:
             raise TypeError(f"Can only add ints and {type(self).__name__}. Instead was {other}")
         return self._with_x(self.x + other)
 
-    def __sub__(self: TSelf, other: int) -> TSelf:
+    def __sub__(self: TSelf, other: Union[int, TSelf]) -> TSelf:
         if isinstance(other, _BaseLineQid):
             if self.dimension != other.dimension:
                 raise TypeError(