Return a single results object instead of always a list - IonQ

cirq-ionq/cirq_ionq/job.py

@@ -16,7 +16,7 @@
 import json
 import time
 import warnings
-from typing import Dict, Optional, Sequence, TYPE_CHECKING, Union
+from typing import Dict, List, Optional, Sequence, TYPE_CHECKING, Union
 
 import cirq
 from cirq._doc import document
@@ -195,7 +195,12 @@ def results(
         polling_seconds: int = 1,
         sharpen: Optional[bool] = None,
         extra_query_params: Optional[dict] = None,
-    ) -> Union[list[results.QPUResult], list[results.SimulatorResult]]:
+    ) -> Union[
+        results.QPUResult,
+        results.SimulatorResult,
+        List[results.QPUResult],
+        List[results.SimulatorResult],
+    ]:
         """Polls the IonQ api for results.
 
         Args:
@@ -242,11 +247,10 @@ def results(
             job_id=self.job_id(), sharpen=sharpen, extra_query_params=extra_query_params
         )
 
+        # is this a batch run (dict‑of‑dicts) or a single circuit?
         some_inner_value = next(iter(backend_results.values()))
-        if isinstance(some_inner_value, dict):
-            histograms = backend_results.values()
-        else:
-            histograms = [backend_results]
+        is_batch = isinstance(some_inner_value, dict)
+        histograms = list(backend_results.values()) if is_batch else [backend_results]
 
         # IonQ returns results in little endian, but
         # Cirq prefers to use big endian, so we convert.
@@ -267,7 +271,11 @@ def results(
                         measurement_dict=self.measurement_dict(circuit_index=circuit_index),
                     )
                 )
-            return big_endian_results_qpu
+            return (
+                big_endian_results_qpu
+                if len(big_endian_results_qpu) > 1
+                else big_endian_results_qpu[0]
+            )
         else:
             big_endian_results_sim: list[results.SimulatorResult] = []
             for circuit_index, histogram in enumerate(histograms):
@@ -283,7 +291,11 @@ def results(
                         repetitions=self.repetitions(),
                     )
                 )
-            return big_endian_results_sim
+            return (
+                big_endian_results_sim
+                if len(big_endian_results_sim) > 1
+                else big_endian_results_sim[0]
+            )
 
     def cancel(self):
         """Cancel the given job.

cirq-ionq/cirq_ionq/job_test.py

@@ -97,7 +97,7 @@ def test_job_results_qpu():
         assert "foo" in str(w[0].message)
         assert "bar" in str(w[1].message)
     expected = ionq.QPUResult({0: 600, 1: 400}, 2, {'a': [0, 1]})
-    assert results[0] == expected
+    assert results == expected
 
 
 def test_batch_job_results_qpu():
@@ -146,7 +146,7 @@ def test_job_results_rounding_qpu():
     job = ionq.Job(mock_client, job_dict)
     expected = ionq.QPUResult({0: 3, 1: 4997}, 2, {'a': [0, 1]})
     results = job.results()
-    assert results[0] == expected
+    assert results == expected
 
 
 def test_job_results_failed():
@@ -177,7 +177,7 @@ def test_job_results_qpu_endianness():
     }
     job = ionq.Job(mock_client, job_dict)
     results = job.results()
-    assert results[0] == ionq.QPUResult({0: 600, 2: 400}, 2, measurement_dict={})
+    assert results == ionq.QPUResult({0: 600, 2: 400}, 2, measurement_dict={})
 
 
 def test_batch_job_results_qpu_endianness():
@@ -198,7 +198,7 @@ def test_batch_job_results_qpu_endianness():
     }
     job = ionq.Job(mock_client, job_dict)
     results = job.results()
-    assert results[0] == ionq.QPUResult({0: 600, 2: 400}, 2, measurement_dict={'a': [0, 1]})
+    assert results == ionq.QPUResult({0: 600, 2: 400}, 2, measurement_dict={'a': [0, 1]})
 
 
 def test_job_results_qpu_target_endianness():
@@ -214,7 +214,7 @@ def test_job_results_qpu_target_endianness():
     }
     job = ionq.Job(mock_client, job_dict)
     results = job.results()
-    assert results[0] == ionq.QPUResult({0: 600, 2: 400}, 2, measurement_dict={})
+    assert results == ionq.QPUResult({0: 600, 2: 400}, 2, measurement_dict={})
 
 
 def test_batch_job_results_qpu_target_endianness():
@@ -236,7 +236,7 @@ def test_batch_job_results_qpu_target_endianness():
     }
     job = ionq.Job(mock_client, job_dict)
     results = job.results()
-    assert results[0] == ionq.QPUResult({0: 600, 2: 400}, 2, measurement_dict={'a': [0, 1]})
+    assert results == ionq.QPUResult({0: 600, 2: 400}, 2, measurement_dict={'a': [0, 1]})
 
 
 @mock.patch('time.sleep', return_value=None)
@@ -254,7 +254,7 @@ def test_job_results_poll(mock_sleep):
     mock_client.get_results.return_value = {'0': '0.6', '1': '0.4'}
     job = ionq.Job(mock_client, ready_job)
     results = job.results(polling_seconds=0)
-    assert results[0] == ionq.QPUResult({0: 600, 1: 400}, 1, measurement_dict={})
+    assert results == ionq.QPUResult({0: 600, 1: 400}, 1, measurement_dict={})
     mock_sleep.assert_called_once()
 
 
@@ -292,7 +292,7 @@ def test_job_results_simulator():
     }
     job = ionq.Job(mock_client, job_dict)
     results = job.results()
-    assert results[0] == ionq.SimulatorResult({0: 0.6, 1: 0.4}, 1, {}, 100)
+    assert results == ionq.SimulatorResult({0: 0.6, 1: 0.4}, 1, {}, 100)
 
 
 def test_batch_job_results_simulator():
@@ -334,7 +334,7 @@ def test_job_results_simulator_endianness():
     }
     job = ionq.Job(mock_client, job_dict)
     results = job.results()
-    assert results[0] == ionq.SimulatorResult({0: 0.6, 2: 0.4}, 2, {}, 100)
+    assert results == ionq.SimulatorResult({0: 0.6, 2: 0.4}, 2, {}, 100)
 
 
 def test_batch_job_results_simulator_endianness():
@@ -355,7 +355,7 @@ def test_batch_job_results_simulator_endianness():
     }
     job = ionq.Job(mock_client, job_dict)
     results = job.results()
-    assert results[0] == ionq.SimulatorResult({0: 0.6, 2: 0.4}, 2, {'a': [0, 1]}, 1000)
+    assert results == ionq.SimulatorResult({0: 0.6, 2: 0.4}, 2, {'a': [0, 1]}, 1000)
 
 
 def test_job_sharpen_results():
@@ -370,7 +370,7 @@ def test_job_sharpen_results():
     }
     job = ionq.Job(mock_client, job_dict)
     results = job.results(sharpen=False)
-    assert results[0] == ionq.SimulatorResult({0: 60, 1: 40}, 1, {}, 100)
+    assert results == ionq.SimulatorResult({0: 60, 1: 40}, 1, {}, 100)
 
 
 def test_job_cancel():

cirq-ionq/cirq_ionq/sampler.py

@@ -13,8 +13,7 @@
 # limitations under the License.
 """A `cirq.Sampler` implementation for the IonQ API."""
 
-import itertools
-from typing import Optional, Sequence, TYPE_CHECKING
+from typing import Optional, Sequence, TYPE_CHECKING, Union
 
 import cirq
 from cirq_ionq import results
@@ -88,8 +87,8 @@ def run_sweep(
             repetitions: The number of times to sample.
 
         Returns:
-            Either a list of `cirq_ionq.QPUResult` or a list of `cirq_ionq.SimulatorResult`
-            depending on whether the job was running on an actual quantum processor or a simulator.
+            Either a single scalar or list of `cirq_ionq.QPUResult` or `cirq_ionq.SimulatorResult`
+            depending on whether the job or jobs ran on an actual quantum processor or a simulator.
         """
         resolvers = [r for r in cirq.to_resolvers(params)]
         jobs = [
@@ -100,11 +99,16 @@ def run_sweep(
             )
             for resolver in resolvers
         ]
+        # collect results
         if self._timeout_seconds is not None:
-            job_results = [job.results(timeout_seconds=self._timeout_seconds) for job in jobs]
+            raw_results = [j.results(timeout_seconds=self._timeout_seconds) for j in jobs]
         else:
-            job_results = [job.results() for job in jobs]
-        flattened_job_results = list(itertools.chain.from_iterable(job_results))
+            raw_results = [j.results() for j in jobs]
+
+        # each element of `raw_results` might be a single result or a list
+        flattened_job_results: list[Union[results.QPUResult, results.SimulatorResult]] = []
+        for r in raw_results:
+            flattened_job_results.extend(r if isinstance(r, list) else [r])
         cirq_results = []
         for result, params in zip(flattened_job_results, resolvers):
             if isinstance(result, results.QPUResult):

cirq-ionq/cirq_ionq/service.py

@@ -15,6 +15,7 @@
 
 import datetime
 import os
+from collections.abc import Iterable
 from typing import List, Optional, Sequence
 
 import cirq
@@ -124,21 +125,27 @@ def run(
             A `cirq.Result` for running the circuit.
         """
         resolved_circuit = cirq.resolve_parameters(circuit, param_resolver)
-        job_results = self.create_job(
+        job_out = self.create_job(
             circuit=resolved_circuit,
             repetitions=repetitions,
             name=name,
             target=target,
             error_mitigation=error_mitigation,
             extra_query_params=extra_query_params,
         ).results(sharpen=sharpen)
-        if isinstance(job_results[0], results.QPUResult):
-            return job_results[0].to_cirq_result(params=cirq.ParamResolver(param_resolver))
-        if isinstance(job_results[0], results.SimulatorResult):
-            return job_results[0].to_cirq_result(
-                params=cirq.ParamResolver(param_resolver), seed=seed
-            )
-        raise NotImplementedError(f"Unrecognized job result type '{type(job_results[0])}'.")
+
+        # `create_job()` always submits a single circuit, so the API either gives us:
+        #   - a QPUResult / SimulatorResult, or
+        #   - a list of length‑1 (the batch logic in Job.results still wraps it in a list).
+        #   In the latter case we unwrap it here.
+        if isinstance(job_out, list):
+            job_out = job_out[0]
+
+        if isinstance(job_out, results.QPUResult):
+            return job_out.to_cirq_result(params=cirq.ParamResolver(param_resolver))
+        if isinstance(job_out, results.SimulatorResult):
+            return job_out.to_cirq_result(params=cirq.ParamResolver(param_resolver), seed=seed)
+        raise NotImplementedError(f"Unrecognized job result type '{type(job_out)}'.")
 
     def run_batch(
         self,
@@ -186,6 +193,10 @@ def run_batch(
             error_mitigation=error_mitigation,
             extra_query_params=extra_query_params,
         ).results(sharpen=sharpen)
+        assert isinstance(job_results, Iterable), (
+            "Expected job results to be iterable, but got type "
+            f"{type(job_results)}. This is a bug in the IonQ API."
+        )
 
         cirq_results = []
         for job_result in job_results:

cirq-ionq/cirq_ionq/service_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import datetime
+import json
 import os
 from unittest import mock
 
@@ -296,3 +297,77 @@ def test_service_remote_host_default():
 def test_service_remote_host_from_env_var_cirq_ionq_precedence():
     service = ionq.Service(api_key='tomyheart')
     assert service.remote_host == 'http://example.com'
+
+
+def test_service_run_unwraps_single_result_list():
+    """`Service.run` should unwrap `[result]` to `result`."""
+    # set up a real Service object (we'll monkey‑patch its create_job)
+    service = ionq.Service(remote_host="http://example.com", api_key="key")
+
+    # simple 1‑qubit circuit
+    q = cirq.LineQubit(0)
+    circuit = cirq.Circuit(cirq.X(q), cirq.measure(q, key="m"))
+
+    # fabricate a QPUResult and wrap it in a list to mimic an erroneous behavior
+    qpu_result = ionq.QPUResult(counts={1: 1}, num_qubits=1, measurement_dict={"m": [0]})
+    mock_job = mock.MagicMock()
+    mock_job.results.return_value = [qpu_result]  # <- list of length‑1
+
+    # monkey‑patch create_job so Service.run sees our mock_job
+    with mock.patch.object(service, "create_job", return_value=mock_job):
+        out = service.run(circuit=circuit, repetitions=1, target="qpu")
+
+    # expected Cirq result after unwrapping and conversion
+    expected = qpu_result.to_cirq_result(params=cirq.ParamResolver({}))
+
+    assert out == expected
+    mock_job.results.assert_called_once()
+
+
+@pytest.mark.parametrize("target", ["qpu", "simulator"])
+def test_run_batch_preserves_order(target):
+    """``Service.run_batch`` must return results in the same order as the
+    input ``circuits`` list, regardless of how the IonQ API happens to order
+    its per‑circuit results.
+    """
+
+    # Service with a fully mocked HTTP client.
+    service = ionq.Service(remote_host="http://example.com", api_key="key")
+    client = mock.MagicMock()
+    service._client = client
+
+    # Three trivial 1‑qubit circuits, each measuring under a unique key.
+    keys = ["a", "b", "c"]
+    q = cirq.LineQubit(0)
+    circuits = [cirq.Circuit(cirq.measure(q, key=k)) for k in keys]
+
+    client.create_job.return_value = {"id": "job_id", "status": "ready"}
+
+    client.get_job.return_value = {
+        "id": "job_id",
+        "status": "completed",
+        "target": target,
+        "qubits": "1",
+        "metadata": {
+            "shots": "1",
+            "measurements": json.dumps([{"measurement0": f"{k}\u001f0"} for k in keys]),
+            "qubit_numbers": json.dumps([1, 1, 1]),
+        },
+    }
+
+    # Intentionally scramble the order returned by the API: b, a, c.
+    client.get_results.return_value = {
+        "res_b": {"0": "1"},
+        "res_a": {"0": "1"},
+        "res_c": {"0": "1"},
+    }
+
+    results = service.run_batch(circuits, repetitions=1, target=target)
+
+    # The order of measurement keys in the results should match the input
+    # circuit order exactly (a, b, c).
+    assert [next(iter(r.measurements)) for r in results] == keys
+
+    # Smoke‑test on the mocked client usage.
+    client.create_job.assert_called_once()
+    client.get_results.assert_called_once()