diff --git a/frappy/datatypes.py b/frappy/datatypes.py index 16168d9d..e2bf0a86 100644 --- a/frappy/datatypes.py +++ b/frappy/datatypes.py @@ -31,7 +31,7 @@ from base64 import b64decode, b64encode from frappy.errors import ConfigError, ProgrammingError, \ RangeError, WrongTypeError from frappy.lib import clamp, generalConfig -from frappy.lib.enum import Enum +from frappy.lib.enum import Enum, EnumMember from frappy.properties import HasProperties, Property generalConfig.set_default('lazy_number_validation', False) @@ -524,7 +524,7 @@ class EnumType(DataType): try: return self._enum[value] except (KeyError, TypeError): # TypeError will be raised when value is not hashable - if isinstance(value, (int, str)): + if isinstance(value, (int, str, EnumMember)): raise RangeError(f'{shortrepr(value)} is not a member of enum {self._enum!r}') from None raise WrongTypeError(f'{shortrepr(value)} must be either int or str for an enum value') from None diff --git a/frappy_psi/magfield.py b/frappy_psi/magfield.py index b93506e7..145a13fb 100644 --- a/frappy_psi/magfield.py +++ b/frappy_psi/magfield.py @@ -33,11 +33,15 @@ Mode = Enum( DRIVEN=50, ) -Status = Enum( +SimpleStatus = Enum( Drivable.Status, + RAMPING=370, +) + +Status = Enum( + SimpleStatus, PREPARED=150, PREPARING=340, - RAMPING=370, STABILIZING=380, FINALIZING=390, ) @@ -48,6 +52,7 @@ ON = 1 class SimpleMagfield(HasStates, Drivable): value = Parameter('magnetic field', datatype=FloatRange(unit='T')) + status = Parameter(datatype=StatusType(SimpleStatus)) target_min = Limit() target_max = Limit() ramp = Parameter(