Coverage for rta_reconstruction/utils/rta_argparse.py: 80%

1import argparse 

2from enum import Enum 

3from typing import Any 

4 

5 

6class EnumAction(argparse.Action): 

7 """Argparse action for handling Enums 

8 

9 Allows to show help with the enum values as choices, while returning an Enum value 

10 once the argument is parsed. 

11 

12 The implementation is copied from this answer on slackoverflow: https://stackoverflow.com/a/60750535 

13 

14 Examples 

15 -------- 

16 >>> class MyEnum(Enum): 

17 >>> Foo = "foo" 

18 >>> Bar = "bar" 

19 >>> parser = argparse.ArgumentParser() 

20 >>> parser.add_argument('--my_enum', type=MyEnum, action=EnumAction) 

21 >>> parser.print_help() # prints --my_enum {foo,bar} 

22 """ 

23 

24 def __init__(self, **kwargs): 

25 # Pop off the enum type from kwargs 

26 enum_type = kwargs.pop("type", None) 

27 

28 # Ensure an Enum subclass is provided 

29 if enum_type is None: 29 ↛ 30line 29 didn't jump to line 30 because the condition on line 29 was never true

30 raise ValueError( 

31 "Argument `type` must be assigned to an Enum class when using EnumAction. It is not set." 

32 ) 

33 if not issubclass(enum_type, Enum): 33 ↛ 34line 33 didn't jump to line 34 because the condition on line 33 was never true

34 raise TypeError( 

35 "Argument `type` must be an Enum class when the action is `EnumAction`, got type {}".format( 

36 str(type(enum_type)) 

37 ) 

38 ) 

39 

40 # Generate choices from the Enum values 

41 kwargs.setdefault("choices", tuple(e.value for e in enum_type)) 

42 

43 # init the action 

44 super(EnumAction, self).__init__(**kwargs) 

45 

46 # keep the enum type to generate Enum instances from their values when called 

47 self.enum = enum_type 

48 

49 def __call__( 

50 self, parser: argparse.ArgumentParser, namespace: argparse.Namespace, values: Any, option_string: str = None 

51 ): 

52 """Convert value back into a `self._enum` instance 

53 

54 Parameters 

55 ---------- 

56 parser : argparse.ArgumentParser 

57 The parser instance invoking this action 

58 namespace : argparse.Namespace 

59 The namespace object that will be returned by the parser. 

60 values : str 

61 The command line argument string 

62 option_string : str, optional 

63 The option string that was used to invoke this action, by default None 

64 """ 

65 # Convert the value back to an enum instance and set it in the namespace 

66 value = self.enum(values) 

67 setattr(namespace, self.dest, value)