Merge pull request #489 from pepkit/dev_pydantic_arguments_fix

Workaround for shortform arguments

looper/cli_pydantic.py

@@ -30,7 +30,14 @@
 from divvy import select_divvy_config
 
 from . import __version__
-from .command_models.commands import SUPPORTED_COMMANDS, TopLevelParser
+
+from .command_models.arguments import ArgumentEnum
+
+from .command_models.commands import (
+    SUPPORTED_COMMANDS,
+    TopLevelParser,
+    add_short_arguments,
+)
 from .const import *
 from .divvy import DEFAULT_COMPUTE_RESOURCES_NAME, select_divvy_config
 from .exceptions import *
@@ -317,10 +324,14 @@ def main(test_args=None) -> None:
         description="Looper Pydantic Argument Parser",
         add_help=True,
     )
+
+    parser = add_short_arguments(parser, ArgumentEnum)
+
     if test_args:
         args = parser.parse_typed_args(args=test_args)
     else:
         args = parser.parse_typed_args()
+
     return run_looper(args, parser, test_args=test_args)
 
 

looper/command_models/arguments.py

@@ -33,9 +33,13 @@ class Argument(pydantic.fields.FieldInfo):
         `FieldInfo`. These are passed along as they are.
     """
 
-    def __init__(self, name: str, default: Any, description: str, **kwargs) -> None:
+    def __init__(
+        self, name: str, default: Any, description: str, alias: str = None, **kwargs
+    ) -> None:
         self._name = name
-        super().__init__(default=default, description=description, **kwargs)
+        super().__init__(
+            default=default, description=description, alias=alias, **kwargs
+        )
         self._validate()
 
     @property
@@ -77,11 +81,13 @@ class ArgumentEnum(enum.Enum):
 
     IGNORE_FLAGS = Argument(
         name="ignore_flags",
+        alias="-i",
         default=(bool, False),
         description="Ignore run status flags",
     )
     FORCE_YES = Argument(
         name="force_yes",
+        alias="-f",
         default=(bool, False),
         description="Provide upfront confirmation of destruction intent, to skip console query. Default=False",
     )
@@ -100,48 +106,61 @@ class ArgumentEnum(enum.Enum):
 
     FLAGS = Argument(
         name="flags",
+        alias="-f",
         default=(List, []),
         description="Only check samples based on these status flags.",
     )
 
     TIME_DELAY = Argument(
         name="time_delay",
+        alias="-t",
         default=(int, 0),
         description="Time delay in seconds between job submissions (min: 0, max: 30)",
     )
     DRY_RUN = Argument(
-        name="dry_run", default=(bool, False), description="Don't actually submit jobs"
+        name="dry_run",
+        alias="-d",
+        default=(bool, False),
+        description="Don't actually submit jobs",
     )
     COMMAND_EXTRA = Argument(
         name="command_extra",
+        alias="-x",
         default=(str, ""),
         description="String to append to every command",
     )
     COMMAND_EXTRA_OVERRIDE = Argument(
         name="command_extra_override",
+        alias="-y",
         default=(str, ""),
         description="Same as command-extra, but overrides values in PEP",
     )
     LUMP = Argument(
         name="lump",
+        alias="-u",
         default=(float, None),
         description="Total input file size (GB) to batch into one job",
     )
     LUMPN = Argument(
         name="lump_n",
+        alias="-n",
         default=(int, None),
         description="Number of commands to batch into one job",
     )
     LUMPJ = Argument(
         name="lump_j",
+        alias="-j",
         default=(int, None),
         description="Lump samples into number of jobs.",
     )
     LIMIT = Argument(
-        name="limit", default=(int, None), description="Limit to n samples"
+        name="limit", alias="-l", default=(int, None), description="Limit to n samples"
     )
     SKIP = Argument(
-        name="skip", default=(int, None), description="Skip samples by numerical index"
+        name="skip",
+        alias="-k",
+        default=(int, None),
+        description="Skip samples by numerical index",
     )
     CONFIG_FILE = Argument(
         name="config_file",
@@ -165,16 +184,19 @@ class ArgumentEnum(enum.Enum):
     )
     OUTPUT_DIR = Argument(
         name="output_dir",
+        alias="-o",
         default=(str, None),
         description="Output directory",
     )
     SAMPLE_PIPELINE_INTERFACES = Argument(
         name="sample_pipeline_interfaces",
+        alias="-S",
         default=(List, []),
         description="Paths to looper sample config files",
     )
     PROJECT_PIPELINE_INTERFACES = Argument(
         name="project_pipeline_interfaces",
+        alias="-P",
         default=(List, []),
         description="Paths to looper project config files",
     )
@@ -204,6 +226,7 @@ class ArgumentEnum(enum.Enum):
     )
     SKIP_FILE_CHECKS = Argument(
         name="skip_file_checks",
+        alias="-f",
         default=(bool, False),
         description="Do not perform input file checks",
     )
@@ -214,6 +237,7 @@ class ArgumentEnum(enum.Enum):
     )
     COMPUTE = Argument(
         name="compute",
+        alias="-c",
         default=(List, []),
         description="List of key-value pairs (k1=v1)",
     )

looper/command_models/commands.py

@@ -9,6 +9,7 @@
 
 from ..const import MESSAGE_BY_SUBCOMMAND
 from .arguments import Argument, ArgumentEnum
+from pydantic2_argparse import ArgumentParser
 
 
 @dataclass
@@ -234,6 +235,39 @@ def create_model(self) -> Type[pydantic.BaseModel]:
 InitPifaceParserModel = InitPifaceParser.create_model()
 
 
+def add_short_arguments(
+    parser: ArgumentParser, argument_enums: Type[ArgumentEnum]
+) -> ArgumentParser:
+    """
+    This function takes a parser object created under pydantic argparse and adds the short arguments AFTER the initial creation.
+    This is a workaround as pydantic-argparse does not currently support this during initial parser creation.
+
+    :param ArgumentParser parser: parser before adding short arguments
+    :param Type[ArgumentEnum] argument_enums:  enumeration of arguments that contain names and aliases
+    :return ArgumentParser parser: parser after short arguments have been added
+    """
+
+    for cmd in parser._subcommands.choices.keys():
+
+        for argument_enum in list(argument_enums):
+            # First check there is an alias for the argument otherwise skip
+            if argument_enum.value.alias:
+                short_key = argument_enum.value.alias
+                long_key = "--" + argument_enum.value.name.replace(
+                    "_", "-"
+                )  # We must do this because the ArgumentEnum names are transformed during parser creation
+                if long_key in parser._subcommands.choices[cmd]._option_string_actions:
+                    argument = parser._subcommands.choices[cmd]._option_string_actions[
+                        long_key
+                    ]
+                    argument.option_strings = (short_key, long_key)
+                    parser._subcommands.choices[cmd]._option_string_actions[
+                        short_key
+                    ] = argument
+
+    return parser
+
+
 SUPPORTED_COMMANDS = [
     RunParser,
     RerunParser,

tests/smoketests/test_run.py

@@ -23,6 +23,29 @@ def test_cli(prep_temp_pep):
         raise pytest.fail("DID RAISE {0}".format(Exception))
 
 
+def test_cli_shortform(prep_temp_pep):
+    tp = prep_temp_pep
+
+    x = ["run", "--looper-config", tp, "-d"]
+    try:
+        main(test_args=x)
+    except Exception:
+        raise pytest.fail("DID RAISE {0}".format(Exception))
+
+    x = ["run", "--looper-config", tp, "-d", "-l", "2"]
+    try:
+        main(test_args=x)
+    except Exception:
+        raise pytest.fail("DID RAISE {0}".format(Exception))
+
+    tp = prep_temp_pep
+    x = ["run", "--looper-config", tp, "-d", "-n", "2"]
+    try:
+        main(test_args=x)
+    except Exception:
+        raise pytest.fail("DID RAISE {0}".format(Exception))
+
+
 def test_running_csv_pep(prep_temp_pep_csv):
     tp = prep_temp_pep_csv