1515import argparse
1616import warnings
1717
18+ from enum import StrEnum
1819from test .support import os_helper , captured_stderr
1920from unittest import mock
2021
@@ -952,6 +953,35 @@ class TestDisallowLongAbbreviationAllowsShortGroupingPrefix(ParserTestCase):
952953 ]
953954
954955
956+ class TestStrEnumChoices (TestCase ):
957+ class Color (StrEnum ):
958+ RED = "red"
959+ GREEN = "green"
960+ BLUE = "blue"
961+
962+ def test_metavar_formatter_with_strenum (self ):
963+ parser = argparse .ArgumentParser ()
964+ parser .add_argument ('--color' , choices = self .Color )
965+ args = parser .parse_args (['--color' , 'red' ])
966+ self .assertEqual (args .color , self .Color .RED )
967+
968+ def test_expand_help_with_strenum (self ):
969+ parser = argparse .ArgumentParser ()
970+ parser .add_argument ('--color' , choices = self .Color , help = 'Choose a color' )
971+ help_output = parser .format_help ()
972+ self .assertIn ('[--color {red,green,blue}]' , help_output )
973+ self .assertIn (' --color {red,green,blue}' , help_output )
974+
975+ def test_check_value_with_strenum (self ):
976+ parser = argparse .ArgumentParser ()
977+ parser .add_argument ('--color' , choices = self .Color )
978+ self .assertRaisesRegex (
979+ argparse .ArgumentError ,
980+ "invalid choice: yellow \(choose from red, green, blue\)" ,
981+ parser .parse_args ,
982+ ['--color' , 'yellow' ],
983+ )
984+
955985# ================
956986# Positional tests
957987# ================
@@ -2399,7 +2429,7 @@ def test_wrong_argument_subparsers_no_destination_error(self):
23992429 parser .parse_args (('baz' ,))
24002430 self .assertRegex (
24012431 excinfo .exception .stderr ,
2402- r"error: argument {foo,bar}: invalid choice: ' baz' \(choose from ' foo', ' bar' \)\n$"
2432+ r"error: argument {foo,bar}: invalid choice: baz \(choose from foo, bar\)\n$"
24032433 )
24042434
24052435 def test_optional_subparsers (self ):
@@ -6061,7 +6091,7 @@ def test_subparser(self):
60616091 args = parser .parse_args (['x' , '--' , 'run' , '--' , 'a' , '--' , 'b' ])
60626092 self .assertEqual (NS (foo = 'x' , f = None , bar = ['a' , '--' , 'b' ]), args )
60636093 self .assertRaisesRegex (argparse .ArgumentError ,
6064- "invalid choice: '--' " ,
6094+ "invalid choice: -- " ,
60656095 parser .parse_args , ['--' , 'x' , '--' , 'run' , 'a' , 'b' ])
60666096
60676097 def test_subparser_after_multiple_argument_option (self ):
@@ -6075,7 +6105,7 @@ def test_subparser_after_multiple_argument_option(self):
60756105 args = parser .parse_args (['--foo' , 'x' , 'y' , '--' , 'run' , 'a' , 'b' , '-f' , 'c' ])
60766106 self .assertEqual (NS (foo = ['x' , 'y' ], f = 'c' , bar = ['a' , 'b' ]), args )
60776107 self .assertRaisesRegex (argparse .ArgumentError ,
6078- "invalid choice: '--' " ,
6108+ "invalid choice: -- " ,
60796109 parser .parse_args , ['--foo' , 'x' , '--' , '--' , 'run' , 'a' , 'b' ])
60806110
60816111
0 commit comments