Skip to content

Commit 54d4603

Browse files
dstrain115tonybruguier
authored andcommitted
Add iterator support for AbstractJob (quantumlib#5136)
- Add support for next() in AbstractJob - Changes __iter__ to return a generator, which will have support for iteration in for loops, as well as next(). - Also improve the iteration tests to be more thorough and accurate. Fixes: quantumlib#5120
1 parent 693a9d8 commit 54d4603

File tree

2 files changed

+31
-5
lines changed

2 files changed

+31
-5
lines changed

cirq-google/cirq_google/engine/abstract_job.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -181,7 +181,7 @@ def calibration_results(self) -> Sequence['calibration_result.CalibrationResult'
181181
"""
182182

183183
def __iter__(self) -> Iterator[cirq.Result]:
184-
return iter(self.results())
184+
yield from self.results()
185185

186186
# pylint: disable=function-redefined
187187
@overload

cirq-google/cirq_google/engine/abstract_job_test.py

+30-4
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,9 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414
from typing import Dict, List, TYPE_CHECKING
15+
import pytest
16+
import numpy as np
17+
import cirq
1518
from cirq_google.engine.abstract_job import AbstractJob
1619

1720
if TYPE_CHECKING:
@@ -83,17 +86,40 @@ def batched_results(self):
8386
pass
8487

8588
def results(self):
86-
return list(range(5))
89+
return list(
90+
cirq.ResultDict(params={}, measurements={'a': np.asarray([t])}) for t in range(5)
91+
)
8792

8893
def calibration_results(self):
8994
pass
9095

9196

9297
def test_instantiation_and_iteration():
9398
job = MockJob()
99+
100+
# Test length
94101
assert len(job) == 5
95-
assert job[3] == 3
102+
103+
# Test direct indexing
104+
assert job[3].measurements['a'][0] == 3
105+
106+
# Test iterating through for loop
96107
count = 0
97-
for num in job:
98-
assert num == count
108+
for result in job:
109+
assert result.measurements['a'][0] == count
99110
count += 1
111+
112+
# Test iterator using iterator
113+
iterator = iter(job)
114+
result = next(iterator)
115+
assert result.measurements['a'][0] == 0
116+
result = next(iterator)
117+
assert result.measurements['a'][0] == 1
118+
result = next(iterator)
119+
assert result.measurements['a'][0] == 2
120+
result = next(iterator)
121+
assert result.measurements['a'][0] == 3
122+
result = next(iterator)
123+
assert result.measurements['a'][0] == 4
124+
with pytest.raises(StopIteration):
125+
next(iterator)

0 commit comments

Comments
 (0)