fix: loosen Enum type restrictions to strings (#18)

src/braket/_schemas/_version.py

@@ -15,4 +15,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "1.7.2.dev0"
+__version__ = "1.8.0.dev0"

src/braket/device_schema/device_action_properties.py

@@ -12,7 +12,7 @@
 # language governing permissions and limitations under the License.
 
 from enum import Enum
-from typing import List
+from typing import List, Union
 
 from pydantic import BaseModel
 
@@ -32,8 +32,9 @@ class DeviceActionProperties(BaseModel):
     This class defines the actions that can be performed by a device
 
     Attributes:
-        version: List of versions for the actions the device supports
-        actionType: Enum for the action type. Type of the action to be performed.
+        version (List[str]): List of versions for the actions the device supports
+        actionType (Union[DeviceActionType, str]): Enum for the action type.
+             Type of the action to be performed.
 
     Examples:
         >>> import json
@@ -45,4 +46,4 @@ class DeviceActionProperties(BaseModel):
     """
 
     version: List[str]
-    actionType: DeviceActionType
+    actionType: Union[DeviceActionType, str]

src/braket/device_schema/device_capabilities.py

@@ -11,7 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 
-from typing import Dict
+from typing import Dict, Union
 
 from pydantic import BaseModel
 
@@ -25,10 +25,11 @@ class DeviceCapabilities(BaseModel):
     across all the devices
 
     Attributes:
-        service: properties which are common to the Braket service
-        action: Map of the action to its properties
-        deviceParameters: It is the json schema of the deviceParameters for each device. for
-            example the deviceParameter for IonqDeviceCapabilities will be
+        service (DeviceServiceProperties): properties which are common to the Braket service
+        action (Dict[Union[DeviceActionType, str], DeviceActionProperties]): Map of the action
+             to its properties
+        deviceParameters (dict): The json schema of the deviceParameters for each device.
+             For example, the deviceParameter for IonqDeviceCapabilities will be
             IonqDeviceParameters.json_schema()
 
     Examples:
@@ -71,5 +72,5 @@ class DeviceCapabilities(BaseModel):
     """
 
     service: DeviceServiceProperties
-    action: Dict[DeviceActionType, DeviceActionProperties]
+    action: Dict[Union[DeviceActionType, str], DeviceActionProperties]
     deviceParameters: dict

src/braket/device_schema/device_execution_window.py

@@ -13,6 +13,7 @@
 
 from datetime import time
 from enum import Enum
+from typing import Union
 
 from pydantic import BaseModel
 
@@ -52,9 +53,9 @@ class DeviceExecutionWindow(BaseModel):
     This class defines when a device can execute a given task.
 
     Attributes:
-        executionDay: Days of the execution window
-        windowStartHour: UTC 24-hour format of the time when the execution window starts
-        windowEndHour: UTC 24-hour format of the time when the execution window ends
+        executionDay (Union[ExecutionDay, str]): Days of the execution window
+        windowStartHour (time): UTC 24-hour format of the time when the execution window starts
+        windowEndHour (time): UTC 24-hour format of the time when the execution window ends
 
     Examples:
         >>> import json
@@ -67,6 +68,6 @@ class DeviceExecutionWindow(BaseModel):
 
     """
 
-    executionDay: ExecutionDay
+    executionDay: Union[ExecutionDay, str]
     windowStartHour: time
     windowEndHour: time

src/braket/device_schema/dwave/dwave_2000Q_device_level_parameters_v1.py

@@ -11,7 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License
 
-from typing import List, Optional
+from typing import List, Optional, Union
 
 from pydantic import Field
 
@@ -47,8 +47,8 @@ class Dwave2000QDeviceLevelParameters(BraketSchemaBase):
             you must supply the initial state to which the system is set.
         maxResults (Optional[int] = Field(gt=1)): Specifies the maximum number of
             answers returned from the solver.
-        postprocessingType (Optional[PostProcessingType]): Defines what type of postprocessing the
-            system runs online on raw solutions.
+        postprocessingType (Optional[Union[PostProcessingType, str]]): Defines what type
+             of postprocessing the system runs online on raw solutions.
         programmingThermalizationDuration (Optional[int]): Gives the time (in microseconds) to wait
             after programming the QPU for it to cool back to base temperature (i.e.,
             post-programming thermalization time).
@@ -90,7 +90,7 @@ class Dwave2000QDeviceLevelParameters(BraketSchemaBase):
     fluxBiases: Optional[List[float]]
     initialState: Optional[List[int]]
     maxResults: Optional[int] = Field(gt=0)
-    postprocessingType: Optional[PostProcessingType]
+    postprocessingType: Optional[Union[PostProcessingType, str]]
     programmingThermalizationDuration: Optional[int]
     readoutThermalizationDuration: Optional[int]
     reduceIntersampleCorrelation: Optional[bool]

src/braket/device_schema/dwave/dwave_provider_level_parameters_v1.py

@@ -12,7 +12,7 @@
 # language governing permissions and limitations under the License.
 
 from enum import Enum
-from typing import List, Optional
+from typing import List, Optional, Union
 
 from pydantic import Field
 
@@ -62,6 +62,8 @@ class DwaveProviderLevelParameters(BraketSchemaBase):
             you must supply the initial state to which the system is set.
         maxResults (Optional[int] = Field(gt=1)): Specifies the maximum number of
             answers returned from the solver.
+        postprocessingType (Optional[Union[PostProcessingType, str]]): Defines what type
+             of postprocessing the system runs online on raw solutions.
         postprocessingType (Optional[PostProcessingType]): Defines what type of postprocessing the
             system runs online on raw solutions.
         programmingThermalizationDuration (Optional[int]): Gives the time (in microseconds) to wait
@@ -105,7 +107,7 @@ class DwaveProviderLevelParameters(BraketSchemaBase):
     fluxBiases: Optional[List[float]]
     initialState: Optional[List[int]]
     maxResults: Optional[int] = Field(gt=0)
-    postprocessingType: Optional[PostProcessingType]
+    postprocessingType: Optional[Union[PostProcessingType, str]]
     programmingThermalizationDuration: Optional[int]
     readoutThermalizationDuration: Optional[int]
     reduceIntersampleCorrelation: Optional[bool]

src/braket/device_schema/ionq/ionq_device_capabilities_v1.py

@@ -32,8 +32,9 @@ class IonqDeviceCapabilities(BraketSchemaBase, DeviceCapabilities):
     This defines the capabilities of an IonQ device.
 
     Attributes:
-        action(Dict[DeviceActionType, JaqcdDeviceActionProperties]): Actions that an IonQ device
-            can support
+        action(Dict[Union[DeviceActionType, str],
+            Union[OpenQASMDeviceActionProperties, JaqcdDeviceActionProperties]]): Actions that an
+            IonQ device can support
         paradigm(GateModelQpuParadigmProperties): Paradigm properties
         provider(Optional[IonqProviderProperties]): IonQ provider specific properties
 
@@ -104,7 +105,8 @@ class IonqDeviceCapabilities(BraketSchemaBase, DeviceCapabilities):
     )
     braketSchemaHeader: BraketSchemaHeader = Field(default=_PROGRAM_HEADER, const=_PROGRAM_HEADER)
     action: Dict[
-        DeviceActionType, Union[OpenQASMDeviceActionProperties, JaqcdDeviceActionProperties]
+        Union[DeviceActionType, str],
+        Union[OpenQASMDeviceActionProperties, JaqcdDeviceActionProperties],
     ]
     paradigm: GateModelQpuParadigmProperties
     provider: Optional[IonqProviderProperties]

src/braket/device_schema/rigetti/rigetti_device_capabilities_v1.py

@@ -31,7 +31,8 @@ class RigettiDeviceCapabilities(BraketSchemaBase, DeviceCapabilities):
     This defines the capabilities of a Rigetti device.
 
     Attributes:
-        action(Dict[DeviceActionType, JaqcdDeviceActionProperties]): Actions that a
+        action(Dict[Union[DeviceActionType, str],
+            Union[OpenQASMDeviceActionProperties, JaqcdDeviceActionProperties]]): Actions that a
             Rigetti device can support
         paradigm(GateModelQpuParadigmProperties): Paradigm properties of a Rigetti
         provider(Optional[RigettiProviderProperties]): Rigetti provider specific properties
@@ -104,7 +105,8 @@ class RigettiDeviceCapabilities(BraketSchemaBase, DeviceCapabilities):
     )
     braketSchemaHeader: BraketSchemaHeader = Field(default=_PROGRAM_HEADER, const=_PROGRAM_HEADER)
     action: Dict[
-        DeviceActionType, Union[OpenQASMDeviceActionProperties, JaqcdDeviceActionProperties]
+        Union[DeviceActionType, str],
+        Union[OpenQASMDeviceActionProperties, JaqcdDeviceActionProperties],
     ]
     paradigm: GateModelQpuParadigmProperties
     provider: Optional[RigettiProviderProperties]

src/braket/device_schema/simulators/gate_model_simulator_device_capabilities_v1.py

@@ -30,7 +30,8 @@ class GateModelSimulatorDeviceCapabilities(BraketSchemaBase, DeviceCapabilities)
     This defines the capabilities of a simulator device.
 
     Attributes:
-        action (Dict[DeviceActionType, JaqcdDeviceActionProperties]): Actions that a
+        action (Dict[Union[DeviceActionType, str],
+            Union[OpenQASMDeviceActionProperties, JaqcdDeviceActionProperties]]): Actions that a
             gate model simulator device can support
         paradigm (GateModelSimulatorParadigmProperties): Paradigm properties of a simulator
 
@@ -99,6 +100,7 @@ class GateModelSimulatorDeviceCapabilities(BraketSchemaBase, DeviceCapabilities)
     )
     braketSchemaHeader: BraketSchemaHeader = Field(default=_PROGRAM_HEADER, const=_PROGRAM_HEADER)
     action: Dict[
-        DeviceActionType, Union[OpenQASMDeviceActionProperties, JaqcdDeviceActionProperties]
+        Union[DeviceActionType, str],
+        Union[OpenQASMDeviceActionProperties, JaqcdDeviceActionProperties],
     ]
     paradigm: GateModelSimulatorParadigmProperties

src/braket/ir/annealing/problem_v1.py

@@ -12,7 +12,7 @@
 # language governing permissions and limitations under the License.
 
 from enum import Enum
-from typing import Dict
+from typing import Dict, Union
 
 from pydantic import Field, conint
 
@@ -36,7 +36,7 @@ class Problem(BraketSchemaBase):
     Attributes:
         braketSchemaHeader (BraketSchemaHeader): Schema header. Users do not need
             to set this value. Only default is allowed.
-        type (ProblemType): The type of problem; can be either "QUBO" or "ISING"
+        type (Union[ProblemType, str]): The type of problem; can be either "QUBO" or "ISING"
         linear (Dict[int, float]): Linear terms of the model.
         quadratic (Dict[str, float]): Quadratic terms of the model, keyed on comma-separated
             variables as strings
@@ -47,6 +47,6 @@ class Problem(BraketSchemaBase):
 
     _PROBLEM_HEADER = BraketSchemaHeader(name="braket.ir.annealing.problem", version="1")
     braketSchemaHeader: BraketSchemaHeader = Field(default=_PROBLEM_HEADER, const=_PROBLEM_HEADER)
-    type: ProblemType
+    type: Union[ProblemType, str]
     linear: Dict[conint(ge=0), float]
     quadratic: Dict[str, float]

src/braket/jobs_data/persisted_job_data_v1.py

@@ -12,7 +12,7 @@
 # language governing permissions and limitations under the License
 
 from enum import Enum
-from typing import Any, Dict
+from typing import Any, Dict, Union
 
 from pydantic import Field
 
@@ -37,8 +37,8 @@ class PersistedJobData(BraketSchemaBase):
         braketSchemaHeader (BraketSchemaHeader): Schema header. Users do not need
             to set this value.
         dataDictionary (Dict[str, Any]): Dict representing the data to be persisted.
-        dataFormat (PersistedJobDataFormat): Data format used for persisting the values
-            in `dataDictionary`.
+        dataFormat (Union[PersistedJobDataFormat, str]): Data format used for persisting the
+            values in `dataDictionary`.
 
     Examples:
         >>> data_to_persist = {"some_key": "some_value", "more_keys": True}
@@ -54,4 +54,4 @@ class PersistedJobData(BraketSchemaBase):
         default=_PERSISTED_JOB_DATA_HEADER, const=_PERSISTED_JOB_DATA_HEADER
     )
     dataDictionary: Dict[str, Any]
-    dataFormat: PersistedJobDataFormat
+    dataFormat: Union[PersistedJobDataFormat, str]

test/unit_tests/braket/device_schema/test_device_action_properties.py

@@ -33,6 +33,12 @@ def test_valid(valid_input):
     assert result.actionType == "braket.ir.jaqcd.program"
 
 
+def test_valid_str_invalid_enum(valid_input):
+    valid_input["actionType"] = "blah"
+    result = DeviceActionProperties.parse_raw(json.dumps(valid_input))
+    assert result.actionType == "blah"
+
+
 @pytest.mark.xfail(raises=ValidationError)
 def test__missing_actionType(valid_input):
     valid_input.pop("actionType")

test/unit_tests/braket/device_schema/test_device_capabilities.py

@@ -52,6 +52,13 @@ def test_valid(valid_input):
     assert DeviceCapabilities.parse_raw(json.dumps(valid_input))
 
 
+def test_valid_action_str(valid_input):
+    action = valid_input["action"]
+    action["blah"] = action["braket.ir.jaqcd.program"]
+    action.pop("braket.ir.jaqcd.program")
+    assert DeviceCapabilities.parse_raw(json.dumps(valid_input))
+
+
 @pytest.mark.xfail(raises=ValidationError)
 def test_missing_action(valid_input):
     valid_input.pop("action")

test/unit_tests/braket/device_schema/test_device_execution_window.py

@@ -38,9 +38,3 @@ def test_valid(valid_input):
 def test__missing_executionDay(valid_input):
     valid_input.pop("executionDay")
     assert DeviceExecutionWindow.parse_raw(json.dumps(valid_input))
-
-
-@pytest.mark.xfail(raises=ValidationError)
-def test__invalid_executionDay(valid_input):
-    valid_input["executionDay"] = "today"
-    DeviceExecutionWindow.parse_raw(json.dumps(valid_input))