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, Optional, Sequence, TYPE_CHECKING, Union, List
 
 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

@@ -100,11 +100,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 = []
+        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

@@ -124,21 +124,24 @@ 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])}'.")
+
+        # normalise: single‑circuit jobs should deliver one result
+        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,