Browse Source

renamed choice enum class and changed some class method logic

Dimitri Korsch 6 years ago
parent
commit
3aec681b97
1 changed files with 12 additions and 8 deletions
  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__))