diff --git a/dargs/dargs.py b/dargs/dargs.py index f067b01..ed698db 100644 --- a/dargs/dargs.py +++ b/dargs/dargs.py @@ -898,6 +898,14 @@ def dummy_argument(self) -> Argument: def get_choice(self, argdict: dict, path: list[str] | None = None) -> Argument: if self.flag_name in argdict: tag = argdict[self.flag_name] + choices = [*self.choice_dict, *self.choice_alias] + if not isinstance(tag, str): + raise ArgumentTypeError( + path, + f"key `{self.flag_name}` gets wrong value type, requires " + f"but {type(tag).__name__} is given; " + f"expected choices are <{'|'.join(choices)}>", + ) if tag in self.choice_dict: return self.choice_dict[tag] elif tag in self.choice_alias: @@ -906,10 +914,7 @@ def get_choice(self, argdict: dict, path: list[str] | None = None) -> Argument: raise ArgumentValueError( path, f"get invalid choice `{tag}` for flag key `{self.flag_name}`." - + did_you_mean( - tag, - list(self.choice_dict.keys()) + list(self.choice_alias.keys()), - ), + + did_you_mean(tag, choices), ) elif self.optional: return self.choice_dict[self.default_tag] @@ -934,7 +939,11 @@ def _convert_choice_alias( ) -> None: if self.flag_name in argdict: tag = argdict[self.flag_name] - if tag not in self.choice_dict and tag in self.choice_alias: + if ( + isinstance(tag, str) + and tag not in self.choice_dict + and tag in self.choice_alias + ): argdict[self.flag_name] = self.choice_alias[tag] # above are traversing part diff --git a/tests/test_checker.py b/tests/test_checker.py index 4486c7a..4143f36 100644 --- a/tests/test_checker.py +++ b/tests/test_checker.py @@ -165,6 +165,25 @@ def test_sub_repeat_dict(self) -> None: with self.assertRaises(ArgumentTypeError): ca.check(err_dict3) + def test_variant_choice_type(self) -> None: + ca = Argument( + "base", + dict, + sub_variants=[Variant("kind", [Argument("alpha", dict, alias=["a"])])], + ) + for tag in ([], {}, 1, None): + with self.subTest(tag=tag): + with self.assertRaises(ArgumentTypeError) as cm: + ca.normalize_value({"kind": tag}) + self.assertEqual(cm.exception.path, "") + self.assertIn("requires ", str(cm.exception)) + self.assertIn("expected choices are ", str(cm.exception)) + + self.assertEqual(ca.normalize_value({"kind": "a"}), {"kind": "alpha"}) + with self.assertRaises(ArgumentValueError) as cm: + ca.normalize_value({"kind": "alpah"}) + self.assertIn("Did you mean: alpha?", str(cm.exception)) + def test_sub_variants(self) -> None: ca = Argument( "base",