Allow any object with supporting protocols to be the action in act_on (…

…quantumlib#5111)

* Allow any object with supporting protocols to be the action in act_on

* lint

cirq-core/cirq/contrib/quimb/mps_simulator.py

@@ -677,7 +677,7 @@ def apply_op(self, op: 'cirq.Operation', prng: np.random.RandomState):
 
     def _act_on_fallback_(
         self,
-        action: Union['cirq.Operation', 'cirq.Gate'],
+        action: Any,
         qubits: Sequence['cirq.Qid'],
         allow_decompose: bool = True,
     ) -> bool:

cirq-core/cirq/protocols/act_on_protocol.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Union, Sequence
+from typing import Any, Sequence, TYPE_CHECKING, Union
 
 from typing_extensions import Protocol
 
@@ -89,7 +89,7 @@ def _act_on_(
 
 
 def act_on(
-    action: Union['cirq.Operation', 'cirq.Gate'],
+    action: Any,
     args: 'cirq.OperationTarget',
     qubits: Sequence['cirq.Qid'] = None,
     *,

cirq-core/cirq/protocols/act_on_protocol_test.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Tuple, Union, Sequence
+from typing import Any, Sequence, Tuple
 
 import numpy as np
 import pytest
@@ -36,7 +36,7 @@ def copy(self, deep_copy_buffers: bool = True):
 
     def _act_on_fallback_(
         self,
-        action: Union['cirq.Operation', 'cirq.Gate'],
+        action: Any,
         qubits: Sequence['cirq.Qid'],
         allow_decompose: bool = True,
     ):

cirq-core/cirq/sim/act_on_args_container.py

@@ -16,6 +16,7 @@
 import warnings
 from collections import abc
 from typing import (
+    Any,
     Dict,
     Generic,
     Iterator,
@@ -25,7 +26,6 @@
     Sequence,
     Tuple,
     TYPE_CHECKING,
-    Union,
 )
 
 import numpy as np
@@ -120,16 +120,26 @@ def create_merged_state(self) -> TActOnArgs:
 
     def _act_on_fallback_(
         self,
-        action: Union['cirq.Operation', 'cirq.Gate'],
+        action: Any,
         qubits: Sequence['cirq.Qid'],
         allow_decompose: bool = True,
     ) -> bool:
-        gate = action.gate if isinstance(action, ops.Operation) else action
+        gate_opt = (
+            action
+            if isinstance(action, ops.Gate)
+            else action.gate
+            if isinstance(action, ops.Operation)
+            else None
+        )
 
-        if isinstance(gate, ops.IdentityGate):
+        if isinstance(gate_opt, ops.IdentityGate):
             return True
 
-        if isinstance(gate, ops.SwapPowGate) and gate.exponent % 2 == 1 and gate.global_shift == 0:
+        if (
+            isinstance(gate_opt, ops.SwapPowGate)
+            and gate_opt.exponent % 2 == 1
+            and gate_opt.global_shift == 0
+        ):
             q0, q1 = qubits
             args0 = self.args[q0]
             args1 = self.args[q1]
@@ -160,8 +170,10 @@ def _act_on_fallback_(
 
         # Decouple any measurements or resets
         if self.split_untangled_states and (
-            isinstance(gate, ops.ResetChannel)
-            or (isinstance(gate, ops.MeasurementGate) and not op_args.ignore_measurement_results)
+            isinstance(gate_opt, ops.ResetChannel)
+            or (
+                isinstance(gate_opt, ops.MeasurementGate) and not op_args.ignore_measurement_results
+            )
         ):
             for q in qubits:
                 if op_args.allows_factoring:

cirq-core/cirq/sim/act_on_args_container_test.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Dict, List, Optional, Sequence, Union
+from typing import Any, Dict, List, Optional, Sequence
 
 import cirq
 
@@ -35,7 +35,7 @@ def copy(self) -> 'EmptyActOnArgs':  # type: ignore
 
     def _act_on_fallback_(
         self,
-        action: Union['cirq.Operation', 'cirq.Gate'],
+        action: Any,
         qubits: Sequence['cirq.Qid'],
         allow_decompose: bool = True,
     ) -> bool:

cirq-core/cirq/sim/act_on_args_test.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Sequence, Union
+from typing import Any, Sequence
 
 import numpy as np
 import pytest
@@ -32,7 +32,7 @@ def _perform_measurement(self, qubits):
 
     def _act_on_fallback_(
         self,
-        action: Union['cirq.Operation', 'cirq.Gate'],
+        action: Any,
         qubits: Sequence['cirq.Qid'],
         allow_decompose: bool = True,
     ) -> bool:

cirq-core/cirq/sim/act_on_density_matrix_args.py

@@ -331,7 +331,7 @@ def __init__(
 
     def _act_on_fallback_(
         self,
-        action: Union['cirq.Operation', 'cirq.Gate'],
+        action: Any,
         qubits: Sequence['cirq.Qid'],
         allow_decompose: bool = True,
     ) -> bool:

cirq-core/cirq/sim/act_on_state_vector_args.py

@@ -476,7 +476,7 @@ def subspace_index(
 
     def _act_on_fallback_(
         self,
-        action: Union['cirq.Operation', 'cirq.Gate'],
+        action: Any,
         qubits: Sequence['cirq.Qid'],
         allow_decompose: bool = True,
     ) -> bool:

cirq-core/cirq/sim/clifford/act_on_stabilizer_args.py

@@ -72,7 +72,7 @@ def state(self) -> TStabilizerState:
 
     def _act_on_fallback_(
         self,
-        action: Union['cirq.Operation', 'cirq.Gate'],
+        action: Any,
         qubits: Sequence['cirq.Qid'],
         allow_decompose: bool = True,
     ) -> Union[bool, NotImplementedType]:

cirq-core/cirq/sim/operation_target.py

@@ -14,6 +14,7 @@
 """An interface for quantum states as targets for operations."""
 import abc
 from typing import (
+    Any,
     Dict,
     Generic,
     Iterator,
@@ -49,14 +50,14 @@ def create_merged_state(self) -> TActOnArgs:
     @abc.abstractmethod
     def _act_on_fallback_(
         self,
-        action: Union['cirq.Operation', 'cirq.Gate'],
+        action: Any,
         qubits: Sequence['cirq.Qid'],
         allow_decompose: bool = True,
     ) -> Union[bool, NotImplementedType]:
         """Handles the act_on protocol fallback implementation.
 
         Args:
-            action: Either a gate or an operation to act on.
+            action: A gate, operation, or other to act on.
             qubits: The applicable qubits if a gate is passed as the action.
             allow_decompose: Flag to allow decomposition.
 

cirq-core/cirq/sim/simulator_base_test.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import math
-from typing import Any, Dict, List, Sequence, Tuple, Union
+from typing import Any, Dict, List, Sequence, Tuple
 
 import numpy as np
 import pytest
@@ -68,7 +68,7 @@ def __init__(self, state, qubits, classical_data):
 
     def _act_on_fallback_(
         self,
-        action: Union['cirq.Operation', 'cirq.Gate'],
+        action: Any,
         qubits: Sequence['cirq.Qid'],
         allow_decompose: bool = True,
     ) -> bool: