extend arg parser

pytorch_lightning/trainer/trainer.py

@@ -1,7 +1,7 @@
 import inspect
 import os
 import logging as python_logging
-from argparse import ArgumentParser
+from argparse import ArgumentParser, Namespace
 from typing import Union, Optional, List, Dict, Tuple, Iterable, Any
 
 import torch
@@ -663,51 +663,60 @@ def add_argparse_args(cls, parent_parser: ArgumentParser) -> ArgumentParser:
         # TODO: get "help" from docstring :)
         for arg, arg_types, arg_default in (at for at in cls.get_init_arguments_and_types()
                                             if at[0] not in depr_arg_names):
-
-            for allowed_type in (at for at in allowed_types if at in arg_types):
-                if allowed_type is bool:
-                    def allowed_type(x):
+            arg_types = [at for at in allowed_types if at in arg_types]
+            if not arg_types:
+                # skip argument with not supported type
+                continue
+            arg_kwargs = {}
+            if bool in arg_types:
+                arg_kwargs.update(nargs="?")
+                # if the only arg type is bool
+                if len(arg_types) == 1:
+                    # redefine the type for ArgParser needed
+                    def use_type(x):
                         return bool(parsing.strtobool(x))
-
-                    # Bool args with default of True parsed as flags not key value pair
-                    if arg_types == (bool,) and arg_default is False:
-                        parser.add_argument(
-                            f'--{arg}',
-                            action='store_true',
-                            dest=arg,
-                            help='autogenerated by pl.Trainer'
-                        )
-                        continue
-
-                if arg == 'gpus':
-                    allowed_type = Trainer.allowed_type
-                    arg_default = Trainer.arg_default
-
-                parser.add_argument(
-                    f'--{arg}',
-                    default=arg_default,
-                    type=allowed_type,
-                    dest=arg,
-                    help='autogenerated by pl.Trainer'
-                )
-                break
+                else:
+                    # filter out the bool as we need to use more general
+                    use_type = [at for at in arg_types if at is not bool][0]
+            else:
+                use_type = arg_types[0]
+
+            if arg == 'gpus':
+                use_type = Trainer._allowed_type
+                arg_default = Trainer._arg_default
+
+            parser.add_argument(
+                f'--{arg}',
+                dest=arg,
+                default=arg_default,
+                type=use_type,
+                help='autogenerated by pl.Trainer',
+                **arg_kwargs,
+            )
 
         return parser
 
-    def allowed_type(x):
+    def _allowed_type(x) -> Union[int, str]:
         if ',' in x:
             return str(x)
         else:
             return int(x)
 
-    def arg_default(x):
+    def _arg_default(x) -> Union[int, str]:
         if ',' in x:
             return str(x)
         else:
             return int(x)
 
+    @staticmethod
+    def parse_argparse(cli_parser: ArgumentParser) -> Namespace:
+        """Parse CLI arguments, required for custom bool types."""
+        args = vars(cli_parser.parse_args())
+        args = {k: True if v is None else v for k, v in args.items()}
+        return Namespace(**args)
+
     @classmethod
-    def from_argparse_args(cls, args, **kwargs):
+    def from_argparse_args(cls, args, **kwargs) -> 'Trainer':
 
         params = vars(args)
         params.update(**kwargs)

tests/trainer/test_trainer_cli.py

@@ -88,3 +88,22 @@ def _raise():
 
     with pytest.raises(_UnkArgError):
         parser.parse_args(cli_args)
+
+
+# todo: add also testing for "gpus"
+@pytest.mark.parametrize(['cli_args', 'expected'], [
+    pytest.param('--auto_lr_find', {'auto_lr_find': True}),
+    pytest.param('--auto_lr_find any_string', {'auto_lr_find': 'any_string'}),
+    pytest.param('', {'auto_lr_find': False}),
+])
+def test_argparse_args_parsing(cli_args, expected):
+    """Test multi type argument with bool."""
+    cli_args = cli_args.split(' ') if cli_args else []
+    with mock.patch("argparse._sys.argv", ["any.py"] + cli_args):
+        parser = ArgumentParser(add_help=False)
+        parser = Trainer.add_argparse_args(parent_parser=parser)
+        args = Trainer.parse_argparse(parser)
+
+    for k, v in expected.items():
+        assert getattr(args, k) == v
+    assert Trainer.from_argparse_args(args)