Browse Source

renamed choice enum class and changed some class method logic

Dimitri Korsch 7 năm trước cách đây
mục cha
commit
3aec681b97
1 tập tin đã thay đổi với 12 bổ sung8 xóa
  1. 12 8
      cvargparse/utils/enumerations.py

+ 12 - 8
cvargparse/utils/enumerations.py

@@ -15,7 +15,13 @@ class MetaBaseType(EnumMeta):
 		else:
 		else:
 			return super(MetaBaseType, cls).__contains__(item)
 			return super(MetaBaseType, cls).__contains__(item)
 
 
-class BaseType(Enum, metaclass=MetaBaseType):
+	def __getitem__(cls, key):
+		"""
+			Redefines the "[]" operation.
+		"""
+		return cls.as_choices()[key.lower()]
+
+class BaseChoiceType(Enum, metaclass=MetaBaseType):
 	"""
 	"""
 		Enum base type. Can be used to define argument choices.
 		Enum base type. Can be used to define argument choices.
 		It also enables to quickly creat, get and display the defined choices.
 		It also enables to quickly creat, get and display the defined choices.
@@ -27,11 +33,9 @@ class BaseType(Enum, metaclass=MetaBaseType):
 
 
 	@classmethod
 	@classmethod
 	def get(cls, key):
 	def get(cls, key):
-		key = key.lower()
-		choices = cls.as_choices()
-		if key in choices:
-			return choices.get(key)
-		elif cls.Default:
-			return cls.Default
+		if isinstance(key, str):
+			return cls[key] if key in cls else cls.Default
+		elif isinstance(key, cls):
+			return key
 		else:
 		else:
-			raise KeyError('Unknown key and no default')
+			raise ValueError("Unknown optimizer type: \"{}\"".format(key.__class__.__name__))