Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Workaround for shortform arguments #489

Merged
merged 9 commits into from
May 15, 2024
13 changes: 12 additions & 1 deletion looper/cli_pydantic.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 *
Expand Down Expand Up @@ -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)


Expand Down
34 changes: 29 additions & 5 deletions looper/command_models/arguments.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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",
)
Expand All @@ -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",
Expand All @@ -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",
)
Expand Down Expand Up @@ -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",
)
Expand All @@ -214,6 +237,7 @@ class ArgumentEnum(enum.Enum):
)
COMPUTE = Argument(
name="compute",
alias="-c",
default=(List, []),
description="List of key-value pairs (k1=v1)",
)
Expand Down
34 changes: 34 additions & 0 deletions looper/command_models/commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@

from ..const import MESSAGE_BY_SUBCOMMAND
from .arguments import Argument, ArgumentEnum
from pydantic2_argparse import ArgumentParser


@dataclass
Expand Down Expand Up @@ -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,
Expand Down
23 changes: 23 additions & 0 deletions tests/smoketests/test_run.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
Loading