cirq.measure - accept list arguments

cirq-core/cirq/ops/measure_util.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Callable, Iterable, List, Optional, Tuple, TYPE_CHECKING, Union
+from typing import Callable, Iterable, List, overload, Optional, Tuple, TYPE_CHECKING, Union
 
 import numpy as np
 
@@ -81,8 +81,29 @@ def measure_paulistring_terms(
     return [PauliMeasurementGate([pauli_basis[q]], key=key_func(q)).on(q) for q in pauli_basis]
 
 
+# pylint: disable=function-redefined
+
+
+@overload
+def measure(
+    *target: raw_types.Qid,
+    key: Optional[Union[str, 'cirq.MeasurementKey']] = None,
+    invert_mask: Tuple[bool, ...] = (),
+) -> raw_types.Operation:
+    pass
+
+
+@overload
+def measure(
+    *target: Iterable[raw_types.Qid],
+    key: Optional[Union[str, 'cirq.MeasurementKey']] = None,
+    invert_mask: Tuple[bool, ...] = (),
+) -> raw_types.Operation:
+    pass
+
+
 def measure(
-    *target: Union[raw_types.Qid, Iterable[Any]],
+    *target,
     key: Optional[Union[str, 'cirq.MeasurementKey']] = None,
     invert_mask: Tuple[bool, ...] = (),
 ) -> raw_types.Operation:
@@ -92,8 +113,8 @@ def measure(
 
     Args:
         *target: The qubits that the measurement gate should measure.
-            This can be provided as variable-length arguments or
-            a combination of nested iterables.
+            These can be specified as separate function arguments or
+            with a single argument for an iterable of qubits.
         key: The string key of the measurement. If this is None, it defaults
             to a comma-separated list of the target qubits' str values.
         invert_mask: A list of Truthy or Falsey values indicating whether
@@ -106,47 +127,64 @@ def measure(
     Raises:
         ValueError: If the qubits are not instances of Qid.
     """
-    targetlist: List[raw_types.Qid] = []
-    for qubit in _flatten_recursively(target, keep=(bytes, str, np.ndarray)):
+    one_iterable_arg: bool = (
+        len(target) == 1
+        and isinstance(target[0], Iterable)
+        and not isinstance(target[0], (bytes, str, np.ndarray))
+    )
+    targets = tuple(target[0]) if one_iterable_arg else target
+    for qubit in targets:
         if isinstance(qubit, np.ndarray):
             raise ValueError(
                 'measure() was called a numpy ndarray. Perhaps you meant '
                 'to call measure_state_vector on numpy array?'
             )
         elif not isinstance(qubit, raw_types.Qid):
             raise ValueError('measure() was called with type different than Qid.')
-        targetlist.append(qubit)
 
     if key is None:
-        key = _default_measurement_key(targetlist)
-    qid_shape = protocols.qid_shape(targetlist)
-    return MeasurementGate(len(targetlist), key, invert_mask, qid_shape).on(*targetlist)
+        key = _default_measurement_key(targets)
+    qid_shape = protocols.qid_shape(targets)
+    return MeasurementGate(len(targets), key, invert_mask, qid_shape).on(*targets)
+
+
+@overload
+def measure_each(
+    *qubits: raw_types.Qid, key_func: Callable[[raw_types.Qid], str] = str
+) -> List[raw_types.Operation]:
+    pass
 
 
+@overload
 def measure_each(
-    *qubits: 'cirq.Qid', key_func: Callable[[raw_types.Qid], str] = str
+    *qubits: Iterable[raw_types.Qid], key_func: Callable[[raw_types.Qid], str] = str
+) -> List[raw_types.Operation]:
+    pass
+
+
+def measure_each(
+    *qubits, key_func: Callable[[raw_types.Qid], str] = str
 ) -> List[raw_types.Operation]:
     """Returns a list of operations individually measuring the given qubits.
 
     The qubits are measured in the computational basis.
 
     Args:
-        *qubits: The qubits to measure.
+        *qubits: The qubits to measure.  These can be passed as separate
+            function arguments or as a one-argument iterable of qubits.
         key_func: Determines the key of the measurements of each qubit. Takes
             the qubit and returns the key for that qubit. Defaults to str.
 
     Returns:
         A list of operations individually measuring the given qubits.
     """
-    return [MeasurementGate(1, key_func(q), qid_shape=(q.dimension,)).on(q) for q in qubits]
-
-
-def _flatten_recursively(items: Iterable[Any], keep: Optional[Tuple[type, ...]]) -> Iterable[Any]:
-    for item in items:
-        if keep is not None and isinstance(item, keep):
-            yield item
-        elif isinstance(item, Iterable):
-            for subitem in _flatten_recursively(item, keep=keep):
-                yield subitem
-        else:
-            yield item
+    one_iterable_arg: bool = (
+        len(qubits) == 1
+        and isinstance(qubits[0], Iterable)
+        and not isinstance(qubits[0], (bytes, str))
+    )
+    qubitsequence = qubits[0] if one_iterable_arg else qubits
+    return [MeasurementGate(1, key_func(q), qid_shape=(q.dimension,)).on(q) for q in qubitsequence]
+
+
+# pylint: enable=function-redefined

cirq-core/cirq/ops/measure_util_test.py

@@ -32,7 +32,7 @@ def test_measure_qubits():
     assert cirq.measure(a) == cirq.MeasurementGate(num_qubits=1, key='a').on(a)
     assert cirq.measure([a]) == cirq.MeasurementGate(num_qubits=1, key='a').on(a)
     assert cirq.measure(a, b) == cirq.MeasurementGate(num_qubits=2, key='a,b').on(a, b)
-    assert cirq.measure([a, [b]]) == cirq.MeasurementGate(num_qubits=2, key='a,b').on(a, b)
+    assert cirq.measure([a, b]) == cirq.MeasurementGate(num_qubits=2, key='a,b').on(a, b)
     qubit_generator = (q for q in (a, b))
     assert cirq.measure(qubit_generator) == cirq.MeasurementGate(num_qubits=2, key='a,b').on(a, b)
     assert cirq.measure(b, a) == cirq.MeasurementGate(num_qubits=2, key='b,a').on(b, a)
@@ -50,23 +50,28 @@ def test_measure_qubits():
     with pytest.raises(ValueError, match='ndarray'):
         _ = cirq.measure(np.array([1, 0]))
 
-    with pytest.raises(ValueError, match='ndarray'):
-        _ = cirq.measure([a, np.array([1, 0])])
-
     with pytest.raises(ValueError, match='Qid'):
         _ = cirq.measure("bork")
 
     with pytest.raises(ValueError, match='Qid'):
-        _ = cirq.measure([a, "bork"])
+        _ = cirq.measure([a, [b]])
+
+    with pytest.raises(ValueError, match='Qid'):
+        _ = cirq.measure([a], [b])
 
 
 def test_measure_each():
     a = cirq.NamedQubit('a')
     b = cirq.NamedQubit('b')
 
     assert cirq.measure_each() == []
+    assert cirq.measure_each([]) == []
     assert cirq.measure_each(a) == [cirq.measure(a)]
+    assert cirq.measure_each([a]) == [cirq.measure(a)]
     assert cirq.measure_each(a, b) == [cirq.measure(a), cirq.measure(b)]
+    assert cirq.measure_each([a, b]) == [cirq.measure(a), cirq.measure(b)]
+    qubit_generator = (q for q in (a, b))
+    assert cirq.measure_each(qubit_generator) == [cirq.measure(a), cirq.measure(b)]
     assert cirq.measure_each(a.with_dimension(3), b.with_dimension(3)) == [
         cirq.measure(a.with_dimension(3)),
         cirq.measure(b.with_dimension(3)),