Skip to content

Commit

Permalink
Add iterator support for AbstractJob
Browse files Browse the repository at this point in the history
- 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().

Fixes: quantumlib#5120
  • Loading branch information
dstrain115 committed Mar 24, 2022
1 parent 33580bb commit eb86f72
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 5 deletions.
2 changes: 1 addition & 1 deletion cirq-google/cirq_google/engine/abstract_job.py
Original file line number Diff line number Diff line change
Expand Up @@ -181,7 +181,7 @@ def calibration_results(self) -> Sequence['calibration_result.CalibrationResult'
"""

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

# pylint: disable=function-redefined
@overload
Expand Down
34 changes: 30 additions & 4 deletions cirq-google/cirq_google/engine/abstract_job_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Dict, List, TYPE_CHECKING
import pytest
import numpy as np
import cirq
from cirq_google.engine.abstract_job import AbstractJob

if TYPE_CHECKING:
Expand Down Expand Up @@ -83,17 +86,40 @@ def batched_results(self):
pass

def results(self):
return list(range(5))
return list(
cirq.ResultDict(params={}, measurements={'a': np.asarray([t])}) for t in range(5)
)

def calibration_results(self):
pass


def test_instantiation_and_iteration():
job = MockJob()

# Test length
assert len(job) == 5
assert job[3] == 3

# Test direct indexing
assert job[3].measurements['a'][0] == 3

# Test iterating through for loop
count = 0
for num in job:
assert num == count
for result in job:
assert result.measurements['a'][0] == count
count += 1

# Test iterator using iterator
iterator = iter(job)
result = next(iterator)
assert result.measurements['a'][0] == 0
result = next(iterator)
assert result.measurements['a'][0] == 1
result = next(iterator)
assert result.measurements['a'][0] == 2
result = next(iterator)
assert result.measurements['a'][0] == 3
result = next(iterator)
assert result.measurements['a'][0] == 4
with pytest.raises(StopIteration):
next(iterator)

0 comments on commit eb86f72

Please sign in to comment.