Source code for autojob.utils.cli
"""Utilities for CLI functions."""
from collections.abc import Iterable
from contextlib import suppress
import logging
import re
from typing import Any
from typing import ClassVar
import click
from autojob import SETTINGS
from autojob.coordinator.validation import val_to_native
from autojob.utils.parsing import convert
from autojob.utils.schemas import Unset
logger = logging.getLogger(__name__)
[docs]
class MemoryFloat(click.ParamType):
"""A float representing an amount of memory."""
name: ClassVar[str] = "memory float"
[docs]
def convert(self, value: str | float, param, ctx) -> float:
"""Convert a memory specification into bytes."""
if not isinstance(value, str | float):
msg = f"{param.name} should be a string or a float: {value!r}"
self.fail(msg, param=param, ctx=ctx)
with suppress(ValueError):
return float(value)
if match := re.match(
r"(?P<memory>\d+(?:\.\d+)?)(?P<units>(?:K|M|G)(?:B)?)", value
):
memory = float(match.group("memory"))
units = match.group("units")
return convert(memory=memory, from_units=units, to_units="B")
msg = f"'{value}' is not a valid memory specification"
self.fail(msg, param=param, ctx=ctx)
[docs]
def mods_to_dict(_: Any, param: str, value: Iterable[str]) -> dict[str, Any]:
"""Convert an iterable of key-value pairs to a dictionary.
Args:
_: The first argument is ignored but retained for `click`
compatibility.
param: The name of the parameter being set (e.g., `calc_mods` or
`slurm_mods`).
value: An iterable of key-value pairs should exist as a string in the
form "key=value". Note that only those values supported by
`~validation.val_to_native` can be correctly parsed.
Returns:
A dictionary mapping calculator parameter names to their Python values.
"""
if not isinstance(value, Iterable):
msg = f"Something is wrong. {param} should be an iterable: {value!r}"
raise click.BadParameter(message=msg)
mods = {}
for x in value:
parameter, *assignment_value = x.split("=", maxsplit=1)
if not parameter or "=" not in x:
msg = (
"Incorrect parameter format. Each argument specified to "
f"{param} must have the form 'key=value' or 'key='"
)
raise click.BadParameter(message=msg)
if assignment_value and (val := assignment_value[0]):
mods[parameter] = val_to_native(val)
else:
mods[parameter] = Unset
return mods
# * Define a decorator (e.g., @reconstructed) which sets this as the object in
# * the context (e.g., ctx.obj = construct_cli_call(ctx=ctx, allowed=allowed))
[docs]
def construct_cli_call(
allowed: list[str] | None = None,
) -> str:
"""Construct the original CLI call.
Args:
allowed: A list of strings indicating which parameters are to be
considered to reconstruct the CLI call.
Returns:
A string representing the command-line call that would produce the
present behaviour.
"""
logger.debug("Constructing CLI call")
allowed = [] if allowed is None else allowed
# This works so long as the CLI only accepts options
ctx = click.get_current_context()
args: list[str] = []
for param, value in ctx.params.items():
if (
ctx.get_parameter_source(param)
== click.core.ParameterSource.COMMANDLINE
and param in allowed
):
if isinstance(value, list | tuple):
for element in value:
args.append(
f"--{param.rstrip('s').replace('_', '-')}"
f'="{str(element).split("/")[-1]}"'
)
elif isinstance(value, dict):
for key, element in value.items():
args.append(
f"--{param.rstrip('s').replace('_', '-')}"
f'="{key}={element!s}"'
)
# This won't work if flag_value=False
elif isinstance(value, bool) and value:
args.append(f"--{param.replace('_', '-')}")
elif param == "verbosity" and value > 0:
args.append(f"-{'v' * value}")
else:
args.append(f'--{param.replace("_", "-")}="{value}"')
command = f"{ctx.command.name} {' '.join(args)}".strip()
logger.debug(f"Successfully constructed CLI call: {command}")
return command