Repository: quantumlib/Cirq

Issue:
EngineJob should not just be iterable but also indexable
Currently only implements `__iter__` but think this should also implement `__getitem__`.


Discussion:

Resolving patch:
diff --git a/cirq/google/engine/engine_job.py b/cirq/google/engine/engine_job.py
--- a/cirq/google/engine/engine_job.py
+++ b/cirq/google/engine/engine_job.py
@@ -15,7 +15,8 @@
 
 import time
 
-from typing import Dict, Iterator, List, Optional, Tuple, TYPE_CHECKING
+from typing import Dict, Iterator, List, Optional, overload, Tuple, \
+    TYPE_CHECKING
 
 from cirq import study
 from cirq.google.engine import calibration
@@ -42,6 +43,10 @@ class EngineJob:
     executing on a machine, or it may have entered a terminal state
     (either succeeding or failing).
 
+    `EngineJob`s can be iterated over, returning `TrialResult`s. These
+    `TrialResult`s can also be accessed by index. Note that this will block
+    until the results are returned from the Engine service.
+
     Attributes:
       project_id: A project_id of the parent Google Cloud Project.
       program_id: Unique ID of the program within the parent project.
@@ -365,6 +370,23 @@ def _raise_on_failure(job: quantum.types.QuantumJob) -> None:
     def __iter__(self) -> Iterator[study.TrialResult]:
         return iter(self.results())
 
+    # pylint: disable=function-redefined
+    @overload
+    def __getitem__(self, item: int) -> study.TrialResult:
+        pass
+
+    @overload
+    def __getitem__(self, item: slice) -> List[study.TrialResult]:
+        pass
+
+    def __getitem__(self, item):
+        return self.results()[item]
+
+    # pylint: enable=function-redefined
+
+    def __len__(self) -> int:
+        return len(self.results())
+
     def __str__(self) -> str:
         return (f'EngineJob(project_id=\'{self.project_id}\', '
                 f'program_id=\'{self.program_id}\', job_id=\'{self.job_id}\')')


Test patch:
diff --git a/cirq/google/engine/engine_job_test.py b/cirq/google/engine/engine_job_test.py
--- a/cirq/google/engine/engine_job_test.py
+++ b/cirq/google/engine/engine_job_test.py
@@ -315,59 +315,98 @@ def test_delete(delete_job):
     delete_job.assert_called_once_with('a', 'b', 'steve')
 
 
+RESULTS = qtypes.QuantumResult(result=_to_any(
+    Merge(
+        """
+sweep_results: [{
+        repetitions: 4,
+        parameterized_results: [{
+            params: {
+                assignments: {
+                    key: 'a'
+                    value: 1
+                }
+            },
+            measurement_results: {
+                key: 'q'
+                qubit_measurement_results: [{
+                  qubit: {
+                    id: '1_1'
+                  }
+                  results: '\006'
+                }]
+            }
+        },{
+            params: {
+                assignments: {
+                    key: 'a'
+                    value: 2
+                }
+            },
+            measurement_results: {
+                key: 'q'
+                qubit_measurement_results: [{
+                  qubit: {
+                    id: '1_1'
+                  }
+                  results: '\005'
+                }]
+            }
+        }]
+    }]
+""", v2.result_pb2.Result())))
+
+
 @mock.patch('cirq.google.engine.engine_client.EngineClient.get_job_results')
 def test_results(get_job_results):
     qjob = qtypes.QuantumJob(execution_status=qtypes.ExecutionStatus(
         state=qtypes.ExecutionStatus.State.SUCCESS))
-    results = qtypes.QuantumResult(result=_to_any(
-        Merge(
-            """
-    sweep_results: [{
-            repetitions: 4,
-            parameterized_results: [{
-                params: {
-                    assignments: {
-                        key: 'a'
-                        value: 1
-                    }
-                },
-                measurement_results: {
-                    key: 'q'
-                    qubit_measurement_results: [{
-                      qubit: {
-                        id: '1_1'
-                      }
-                      results: '\006'
-                    }]
-                }
-            },{
-                params: {
-                    assignments: {
-                        key: 'a'
-                        value: 2
-                    }
-                },
-                measurement_results: {
-                    key: 'q'
-                    qubit_measurement_results: [{
-                      qubit: {
-                        id: '1_1'
-                      }
-                      results: '\005'
-                    }]
-                }
-            }]
-        }]
-    """, v2.result_pb2.Result())))
-    get_job_results.return_value = results
+    get_job_results.return_value = RESULTS
 
     job = cg.EngineJob('a', 'b', 'steve', EngineContext(), _job=qjob)
     data = job.results()
+    assert len(data) == 2
     assert str(data[0]) == 'q=0110'
     assert str(data[1]) == 'q=1010'
     get_job_results.assert_called_once_with('a', 'b', 'steve')
 
 
+@mock.patch('cirq.google.engine.engine_client.EngineClient.get_job_results')
+def test_results_iter(get_job_results):
+    qjob = qtypes.QuantumJob(execution_status=qtypes.ExecutionStatus(
+        state=qtypes.ExecutionStatus.State.SUCCESS))
+    get_job_results.return_value = RESULTS
+
+    job = cg.EngineJob('a', 'b', 'steve', EngineContext(), _job=qjob)
+    results = [str(r) for r in job]
+    assert len(results) == 2
+    assert results[0] == 'q=0110'
+    assert results[1] == 'q=1010'
+
+
+@mock.patch('cirq.google.engine.engine_client.EngineClient.get_job_results')
+def test_results_getitem(get_job_results):
+    qjob = qtypes.QuantumJob(execution_status=qtypes.ExecutionStatus(
+        state=qtypes.ExecutionStatus.State.SUCCESS))
+    get_job_results.return_value = RESULTS
+
+    job = cg.EngineJob('a', 'b', 'steve', EngineContext(), _job=qjob)
+    assert str(job[0]) == 'q=0110'
+    assert str(job[1]) == 'q=1010'
+    with pytest.raises(IndexError):
+        _ = job[2]
+
+
+@mock.patch('cirq.google.engine.engine_client.EngineClient.get_job_results')
+def test_results_len(get_job_results):
+    qjob = qtypes.QuantumJob(execution_status=qtypes.ExecutionStatus(
+        state=qtypes.ExecutionStatus.State.SUCCESS))
+    get_job_results.return_value = RESULTS
+
+    job = cg.EngineJob('a', 'b', 'steve', EngineContext(), _job=qjob)
+    assert len(job) == 2
+
+
 @mock.patch('cirq.google.engine.engine_client.EngineClient.get_job')
 @mock.patch('time.sleep', return_value=None)
 def test_timeout(patched_time_sleep, get_job):