Move casts to a shared method

amazon-braket · Aug 17, 2020 · 8d08830 · 8d08830
1 parent 7104e53
commit 8d08830
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 16 deletions.
diff --git a/src/braket/aws/aws_quantum_task.py b/src/braket/aws/aws_quantum_task.py
@@ -426,12 +426,7 @@ def _format_result(result):
 
 @_format_result.register
 def _(result: GateModelTaskResult) -> GateModelQuantumTaskResult:
-    if result.resultTypes:
-        for result_type in result.resultTypes:
-            type = result_type.type.type
-            if type == "amplitude":
-                for state in result_type.value:
-                    result_type.value[state] = complex(*result_type.value[state])
+    GateModelQuantumTaskResult.cast_result_types(result)
     return GateModelQuantumTaskResult.from_object(result)
 
 

diff --git a/src/braket/tasks/gate_model_quantum_task_result.py b/src/braket/tasks/gate_model_quantum_task_result.py
@@ -217,16 +217,7 @@ def from_string(result: str) -> GateModelQuantumTaskResult:
                 in the result dict
         """
         obj = GateModelTaskResult.parse_raw(result)
-        if obj.resultTypes:
-            for result_type in obj.resultTypes:
-                type = result_type.type.type
-                if type == "probability":
-                    result_type.value = np.array(result_type.value)
-                elif type == "statevector":
-                    result_type.value = np.array([complex(*value) for value in result_type.value])
-                elif type == "amplitude":
-                    for state in result_type.value:
-                        result_type.value[state] = complex(*result_type.value[state])
+        GateModelQuantumTaskResult.cast_result_types(obj)
         return GateModelQuantumTaskResult._from_object_internal(obj)
 
     @classmethod
@@ -303,6 +294,26 @@ def _from_dict_internal_simulator_only(cls, result: GateModelTaskResult):
             values=values,
         )
 
+    @staticmethod
+    def cast_result_types(gate_model_task_result: GateModelTaskResult) -> None:
+        """
+        Casts the result types to the types expected by the SDK.
+
+        Args:
+            gate_model_task_result (GateModelTaskResult): GateModelTaskResult representing the
+                results.
+        """
+        if gate_model_task_result.resultTypes:
+            for result_type in gate_model_task_result.resultTypes:
+                type = result_type.type.type
+                if type == "probability":
+                    result_type.value = np.array(result_type.value)
+                elif type == "statevector":
+                    result_type.value = np.array([complex(*value) for value in result_type.value])
+                elif type == "amplitude":
+                    for state in result_type.value:
+                        result_type.value[state] = complex(*result_type.value[state])
+
     @staticmethod
     def _calculate_result_types(
         ir_string: str, measurements: np.ndarray, measured_qubits: List[int]