snk-cli 0.0.1__py3-none-any.whl

snk_cli/dynamic_typer.py

@@ -0,0 +1,249 @@
+import typer
+from typing import List, Callable
+from inspect import signature, Parameter
+from makefun import with_signature
+
+from .options import Option
+import sys
+
+
+class DynamicTyper:
+    app: typer.Typer
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        cls.app = typer.Typer()
+
+    def __call__(self):
+        """
+        Invoke the CLI.
+        Side Effects:
+          Invokes the CLI.
+        Examples:
+          >>> CLI(Path('/path/to/workflow'))()
+        """
+        self.app()
+
+    def register_default_command(self, command: Callable, **command_kwargs) -> None:
+        """
+        Register a default command to the CLI.
+        Args:
+          command (Callable): The command to register.
+        Side Effects:
+          Registers the command to the CLI.
+        Examples:
+          >>> CLI.register_default_command(my_command)
+        """
+        from makefun import with_signature
+        from inspect import signature, Parameter
+
+        command_signature = signature(command)
+        params = list(command_signature.parameters.values())
+        has_ctx = any([p.name == "ctx" for p in params])
+        if not has_ctx:
+            params.insert(
+                0,
+                Parameter(
+                    "ctx",
+                    kind=Parameter.POSITIONAL_OR_KEYWORD,
+                    annotation=typer.Context,
+                ),
+            )
+            command_signature = command_signature.replace(parameters=params)
+
+        @with_signature(command_signature)
+        def wrapper(ctx: typer.Context, *args, **kwargs):
+            if ctx.invoked_subcommand is None:
+                if has_ctx:
+                    return command(ctx, *args, **kwargs)
+                return command(*args, **kwargs)
+
+        self.register_callback(wrapper, invoke_without_command=True, **command_kwargs)
+
+    def register_command(
+        self, command: Callable, dynamic_options=None, **command_kwargs
+    ) -> None:
+        """
+        Register a command to the CLI.
+        Args:
+          command (Callable): The command to register.
+        Side Effects:
+          Registers the command to the CLI.
+        Examples:
+          >>> CLI.register_command(my_command)
+        """
+        if dynamic_options is not None:
+            command = self.add_dynamic_options(command, dynamic_options)
+        if isinstance(command, DynamicTyper):
+            self.app.registered_commands.extend(command.app.registered_commands)
+        else:
+            self.app.command(**command_kwargs)(command)
+
+    def register_callback(self, command: Callable, **command_kwargs) -> None:
+        """
+        Register a callback to the CLI.
+        Args:
+          command (Callable): The callback to register.
+        Side Effects:
+          Registers the callback to the CLI.
+        Examples:
+          >>> CLI.register_callback(my_callback)
+        """
+        self.app.callback(**command_kwargs)(command)
+
+    def register_group(self, group: "DynamicTyper", **command_kwargs) -> None:
+        """
+        Register a subcommand group group to the CLI.
+        Args:
+          group (DynamicTyper): The subcommand group to register.
+        Side Effects:
+          Registers the subcommand group to the CLI.
+        Examples:
+          >>> CLI.register_app(my_group)
+        """
+        self.app.add_typer(group.app, **command_kwargs)
+
+    def _create_cli_parameter(self, option: Option):
+        """
+        Creates a parameter for a CLI option.
+        Args:
+          option (Option): An Option object containing the option's name, type, required status, default value, and help message.
+        Returns:
+          Parameter: A parameter object for the CLI option.
+        Examples:
+          >>> option = Option(name='foo', type='int', required=True, default=0, help='A number')
+          >>> create_cli_parameter(option)
+          Parameter('foo', kind=Parameter.POSITIONAL_OR_KEYWORD, default=typer.Option(..., help='[CONFIG] A number'), annotation=int)
+        """
+        return Parameter(
+            option.name,
+            kind=Parameter.POSITIONAL_OR_KEYWORD,
+            default=typer.Option(
+                ... if option.required else option.default,
+                *[option.flag, option.short_flag] if option.short else [],
+                help=f"{option.help}",
+                rich_help_panel="Workflow Configuration",
+                hidden=option.hidden,
+            ),
+            annotation=option.type,
+        )
+
+    def check_if_option_passed_via_command_line(self, option: Option):
+        """
+        Check if an option is passed via the command line.
+        Args:
+          option (Option): An Option object containing the option's name, type, required status, default value, and help message.
+        Returns:
+          bool: Whether the option is passed via the command line.
+        """
+        if option.flag in sys.argv:
+            return True
+        elif option.type is bool and f"--no-{option.flag[2:]}" in sys.argv:
+            # Check for boolean flags like --foo/--no-foo
+            return True
+        elif option.short and option.short_flag in sys.argv:
+            return True
+        return False
+
+    def add_dynamic_options(self, func: Callable, options: List[Option]):
+        """
+        Function to add dynamic options to a command.
+        Args:
+          command (Callable): The command to which the dynamic options should be added.
+          options (List[dict]): A list of dictionaries containing the option's name, type, required status, default value, and help message.
+        Returns:
+          Callable: A function with the dynamic options added.
+        Examples:
+          >>> my_func = add_dynamic_options_to_function(my_func, [{'name': 'foo', 'type': 'int', 'required': True, 'default': 0, 'help': 'A number'}])
+          >>> my_func
+          <function my_func at 0x7f8f9f9f9f90>
+        """
+        func_sig = signature(func)
+        params = list(func_sig.parameters.values())
+        for op in options[::-1]:
+            params.insert(1, self._create_cli_parameter(op))
+        new_sig = func_sig.replace(parameters=params)
+
+        @with_signature(func_signature=new_sig, func_name=func.__name__)
+        def func_wrapper(*args, **kwargs):
+            """
+            Wraps a function with dynamic options.
+            Args:
+                *args: Variable length argument list.
+                **kwargs: Arbitrary keyword arguments.
+            Returns:
+                Callable: A wrapped function with the dynamic options added.
+            Notes:
+                This function is used in the `add_dynamic_options_to_function` function.
+            """
+            flat_config = None
+
+            if kwargs.get("configfile"):
+                from snakemake import load_configfile
+                from .utils import flatten
+
+                snakemake_config = load_configfile(kwargs["configfile"])
+                flat_config = flatten(snakemake_config)
+
+            for snk_cli_option in options:
+
+                def add_option_to_args():
+                    kwargs["ctx"].args.extend([f"--{snk_cli_option.name}", kwargs[snk_cli_option.name]])
+
+                passed_via_command_line = self.check_if_option_passed_via_command_line(
+                    snk_cli_option
+                )
+
+                if flat_config is None:
+                    # If no config file is provided then all options should be added to the arguments
+                    # later on we will check to see if they differ from any defaults
+                    add_option_to_args()
+                elif passed_via_command_line:
+                    # If an option is passed via the command line if should override the default
+                    add_option_to_args()
+                elif flat_config and snk_cli_option.original_key not in flat_config:
+                    # If a config file is provided and the snk_cli_option key isn't in it, 
+                    # add the snk_cli_option to the arguments
+                    add_option_to_args()
+
+            kwargs = {
+                k: v for k, v in kwargs.items() if k in func_sig.parameters.keys()
+            }
+            return func(*args, **kwargs)
+
+        return func_wrapper
+
+    def error(self, msg, exit=True):
+        """
+        Logs an error message (red) and exits (optional).
+        Args:
+          msg (str): The error message to log.
+          exit (bool): Whether to exit after logging the error message.
+        """
+        typer.secho(msg, fg="red", err=True)
+        if exit:
+            raise typer.Exit(1)
+
+    def success(self, msg):
+        """
+        Logs a success message (green).
+        Args:
+          msg (str): The success message to log.
+        """
+        typer.secho(msg, fg="green")
+
+    def log(self, msg, color="yellow", stderr=True):
+        """
+        Logs a message (yellow).
+        Args:
+          msg (str): The message to log.
+        """
+        typer.secho(msg, fg=color, err=stderr)
+
+    def echo(self, msg):
+        """
+        Prints a message.
+        Args:
+          msg (str): The message to print.
+        """
+        typer.echo(msg)

snk_cli/options/__init__.py

@@ -0,0 +1 @@
+from .option import Option  # noqa: F401

snk_cli/options/option.py

@@ -0,0 +1,18 @@
+from dataclasses import dataclass
+from typing import Any, Optional
+
+
+@dataclass
+class Option:
+    name: str
+    original_key: str
+    default: Any
+    updated: bool
+    help: str
+    type: Any
+    required: bool
+    short: Optional[str]
+    flag: str
+    short_flag: Optional[str]
+    hidden: bool = False
+    from_annotation: bool = False

snk_cli/options/utils.py

@@ -0,0 +1,117 @@
+from typing import List, Any
+from ..config.config import SnkConfig
+from ..utils import get_default_type, flatten
+from .option import Option
+from pathlib import Path
+
+types = {
+    "int": int,
+    "integer": int,
+    "float": float,
+    "str": str,
+    "string": str,
+    "path": Path,
+    "bool": bool,
+    "boolean": bool,
+    "list": List[str],
+    "list[str]": List[str],
+    "list[path]": List[Path],
+    "list[int]": List[int],
+}
+
+def get_keys_from_annotation(annotations):
+    # Get the unique keys from the annotations
+    # preserving the order
+    keys = []
+    for key in annotations:
+        key = ":".join(key.split(":")[:-1])
+        if key not in keys:
+            keys.append(key)
+    return keys
+
+def create_option_from_annotation(
+    annotation_key: str,
+    annotation_values: dict,
+    default_values: dict,
+    from_annotation: bool = False,
+) -> Option:
+    """
+    Create an Option object from a given annotation.
+    Args:
+      annotation_key: The key in the annotations.
+      annotation_values: The dictionary of annotation values.
+      default_values: default value from config.
+    Returns:
+      An Option object.
+    """    
+    config_default = default_values.get(annotation_key, None)
+
+    default = annotation_values.get(f"{annotation_key}:default", config_default)
+    updated = False
+    if config_default is None or default != config_default:
+        updated = True
+    type = annotation_values.get(f"{annotation_key}:type", get_default_type(default))
+    assert (
+        type is not None
+    ), f"Type for {annotation_key} should be one of {', '.join(types.keys())}."
+    annotation_type = types.get(
+        type.lower(), List[str] if "list" in type.lower() else str
+    )
+    name = annotation_values.get(
+        f"{annotation_key}:name", annotation_key.replace(":", "_")
+    ).replace("-", "_")
+    short = annotation_values.get(f"{annotation_key}:short", None)
+    hidden = annotation_values.get(f"{annotation_key}:hidden", False)
+    return Option(
+        name=name,
+        original_key=annotation_key,
+        default=annotation_values.get(f"{annotation_key}:default", default),
+        updated=updated,
+        help=annotation_values.get(f"{annotation_key}:help", ""),
+        type=annotation_type,
+        required=annotation_values.get(f"{annotation_key}:required", False),
+        short=short,
+        flag=f"--{name.replace('_', '-')}",
+        short_flag=f"-{short}" if short else None,
+        hidden=hidden,
+        from_annotation=from_annotation,
+    )
+
+
+def build_dynamic_cli_options(
+    snakemake_config: dict, snk_config: SnkConfig
+) -> List[dict]:
+    """
+    Builds a list of options from a snakemake config and a snk config.
+    Args:
+      snakemake_config (dict): A snakemake config.
+      snk_config (SnkConfig): A snk config.
+    Returns:
+      List[dict]: A list of options.
+    """
+    flat_config = flatten(snakemake_config)
+    flat_annotations = flatten(snk_config.cli)
+    annotation_keys = get_keys_from_annotation(flat_annotations)
+    options = {}
+
+    # For every parameter in the config, create an option from the corresponding annotation
+    for parameter in flat_config:
+        if parameter not in annotation_keys and snk_config.skip_missing:
+            continue
+        options[parameter] = create_option_from_annotation(
+            parameter,
+            flat_annotations,
+            default_values=flat_config,
+        )
+
+    # For every annotation not in config, create an option with default values    
+    for key in annotation_keys:
+        if key not in options:
+            # in annotation but not in config
+            options[key] = create_option_from_annotation(
+                key,
+                flat_annotations,
+                default_values={},
+                from_annotation=True,
+            )
+    return list(options.values())

snk_cli/subcommands/__init__.py

@@ -0,0 +1,5 @@
+from .env import EnvApp
+from .config import ConfigApp
+from .run import RunApp
+from .script import ScriptApp
+from .profile import ProfileApp

snk_cli/subcommands/config.py

@@ -0,0 +1,64 @@
+from typing import List
+import typer
+
+from snk_cli.dynamic_typer import DynamicTyper
+from snk_cli.options.option import Option
+from ..workflow import Workflow
+
+
+class ConfigApp(DynamicTyper):
+    def __init__(self, workflow: Workflow, options: List[Option]):
+        """
+        Initializes the ConfigApp class.
+        Args:
+            workflow (Workflow): The workflow to configure.
+        """
+        self.options = options
+        self.workflow = workflow
+        self.register_command(self.config, help="Show the workflow configuration.")
+
+    def config(
+        self, ctx: typer.Context, pretty: bool = typer.Option(False, "--pretty", "-p")
+    ):
+        """
+        Prints the configuration for the workflow.
+        Args:
+            pretty (bool, optional): Whether to print the configuration in a pretty format. Defaults to False.
+        Returns:
+            None
+        Examples:
+            >>> ConfigApp.show(pretty=True)
+            # Pretty printed configuration
+        """
+        import yaml
+        from collections import defaultdict
+        from rich.console import Console
+        from rich.syntax import Syntax
+        from snk_cli.utils import convert_key_to_snakemake_format
+
+        def deep_update(source, overrides):
+            for key, value in overrides.items():
+                if isinstance(value, dict):
+                    if not isinstance(source.get(key), dict):
+                        # If the existing value is not a dictionary, replace it with one
+                        source[key] = {}
+                    # Now we are sure that source[key] is a dictionary, so we can update it
+                    deep_update(source[key], value)
+                else:
+                    source[key] = value
+            return source
+
+        collapsed_data = defaultdict(dict)
+        config_dict = [
+            convert_key_to_snakemake_format(option.original_key, option.default)
+            for option in self.options
+        ]
+        for d in config_dict:
+            deep_update(collapsed_data, d)
+        yaml_str = yaml.dump(dict(collapsed_data))
+        if pretty:
+            syntax = Syntax(yaml_str, "yaml")
+            console = Console()
+            console.print(syntax)
+        else:
+            typer.echo(yaml_str)

snk_cli/subcommands/env.py

@@ -0,0 +1,211 @@
+import os
+import shutil
+import subprocess
+from pathlib import Path
+import sys
+from typing import List, Optional
+import typer
+
+from snk_cli.dynamic_typer import DynamicTyper
+from .utils import create_snakemake_workflow
+from ..workflow import Workflow
+from rich.console import Console
+from rich.syntax import Syntax
+from snakemake.deployment.conda import Conda, Env, CreateCondaEnvironmentException
+from snk_cli.config.config import get_config_from_workflow_dir
+
+from concurrent.futures import ProcessPoolExecutor
+
+def get_num_cores(default=4):
+    try:
+        return os.cpu_count()
+    except:
+        return default
+
+def create_conda_environment(args):
+    # unpack args
+    env_path_str, snakefile, snakemake_config, configfiles, conda_prefix_dir_str = args
+    # Reconstruct the snakemake_workflow configuration
+    conda_prefix_dir = Path(conda_prefix_dir_str).resolve()
+    snakemake_workflow = create_snakemake_workflow(
+        snakefile,
+        config=snakemake_config,
+        configfiles=configfiles,
+        use_conda=True,
+        conda_prefix=conda_prefix_dir,
+    )
+    env_path = Path(env_path_str).resolve()
+    env = Env(snakemake_workflow, env_file=env_path)
+    try:
+        env.create()
+    except CreateCondaEnvironmentException as e:
+        typer.secho(str(e), fg="red", err=True)
+        return 1
+    return 0
+
+
+class EnvApp(DynamicTyper):
+    def __init__(
+        self,
+        workflow: Workflow,
+        conda_prefix_dir: Path,
+        snakemake_config,
+        snakefile: Path,
+    ):
+        self.workflow = workflow
+        self.conda_prefix_dir = conda_prefix_dir
+        self.snakemake_config = snakemake_config
+        self.snakefile = snakefile
+        self.configfile = get_config_from_workflow_dir(self.workflow.path)
+        self.snakemake_workflow = create_snakemake_workflow(
+            self.snakefile,
+            config=self.snakemake_config,
+            configfiles=[self.configfile] if self.configfile else None,
+            use_conda=True,
+            conda_prefix=self.conda_prefix_dir.resolve(),
+        )
+        self.register_command(self.list, help="List the environments in the workflow.")
+        self.register_command(self.show, help="Show the contents of an environment.")
+        self.register_command(
+            self.run, help="Run a command in one of the workflow environments."
+        )
+        self.register_command(
+            self.activate, help="Activate a workflow conda environment."
+        )
+        self.register_command(self.remove, help="Remove conda environments.")
+        self.register_command(self.create, help="Create conda environments.")
+
+    def list(
+        self,
+        verbose: bool = typer.Option(
+            False, "--verbose", "-v", help="Show conda paths."
+        ),
+    ):
+        from rich.console import Console
+        from rich.table import Table
+
+        table = Table("Name", "CMD", "Env", show_header=True, show_lines=True)
+        for env in self.workflow.environments:
+            conda = Env(self.snakemake_workflow, env_file=env.resolve())
+            # address relative to cwd
+            address = Path(conda.address)
+            if address.exists():
+                address = str(address) if verbose else f"[green]{address.name}[/green]"
+                cmd = f"{self.workflow.name} env activate {env.stem}"
+            else:
+                address = ""
+                cmd = f"{self.workflow.name} env show {env.stem}"
+            table.add_row(
+                env.stem, cmd, address
+            )
+        console = Console()
+        console.print(table)
+
+    def _get_conda_env_path(self, name: str) -> Path:
+        env = [e for e in self.workflow.environments if e.stem == name]
+        if not env:
+            self.error(f"Environment {name} not found!")
+        return env[0]
+
+    def _shellcmd(self, env_address: str, cmd: str) -> str:
+        if sys.platform.lower().startswith("win"):
+            return Conda().shellcmd_win(env_address, cmd)
+        return Conda().shellcmd(env_address, cmd)
+
+    def show(
+        self,
+        name: str = typer.Argument(..., help="The name of the environment."),
+        pretty: bool = typer.Option(
+            False, "--pretty", "-p", help="Pretty print the environment."
+        ),
+    ):
+        env_path = self._get_conda_env_path(name)
+        env_file_text = env_path.read_text()
+        if pretty:
+            syntax = Syntax(env_file_text, "yaml")
+            console = Console()
+            console.print(syntax)
+        else:
+            typer.echo(env_file_text)
+
+    def run(
+        self,
+        name: str = typer.Argument(..., help="The name of the environment."),
+        verbose: bool = typer.Option(
+            False, "--verbose", "-v", help="Print the command to run."
+        ),
+        cmd: List[str] = typer.Argument(..., help="The command to run in environment."),
+    ):
+        env_path = self._get_conda_env_path(name)
+        env = Env(self.snakemake_workflow, env_file=env_path.resolve())
+        env.create()
+        cmd = self._shellcmd(env.address, " ".join(cmd))
+        if verbose:
+            self.log(cmd)
+        user_shell = os.environ.get("SHELL", "/bin/bash")
+        subprocess.run(cmd, shell=True, env=os.environ.copy(), executable=user_shell)
+
+    def remove(
+        self,
+        name: Optional[str] = typer.Argument(None, help="The name of the environment. If not provided, all environments will be deleted."),
+        force: bool = typer.Option(
+            False, "--force", "-f", help="Force deletion of the environments."
+        ),
+    ):
+        if name:
+            env_path = self._get_conda_env_path(name)
+            env = Env(self.snakemake_workflow, env_file=env_path.resolve())
+            path = Path(env.address)
+            if not path.exists():
+                self.error(f"Environment {name} not created!")
+        else:
+            path = self.conda_prefix_dir
+        if force or input(f"Delete {path}? [y/N] ").lower() == "y":
+            shutil.rmtree(path)
+            self.success(f"Deleted {path}!")
+
+    def create(
+            self, 
+            names: Optional[List[str]] = typer.Argument(None, help="The names of the environments to create. If not provided, all environments will be created."),
+            max_workers: int = typer.Option(get_num_cores(), "--workers", "-w", help="Max number of envs to create in parallel.")
+        ):
+        if names:
+            env_paths = [self._get_conda_env_path(name) for name in names]
+        else:
+            env_paths = self.workflow.environments
+        env_args = [
+            (
+                path,
+                self.snakefile,
+                self.snakemake_config,
+                [self.configfile] if self.configfile else None,
+                self.conda_prefix_dir.resolve(),
+            )
+            for path in env_paths 
+        ]
+        with ProcessPoolExecutor(max_workers=max_workers) as executor:
+            status_codes = executor.map(create_conda_environment, env_args)
+        if any(status_codes):
+            self.error("Failed to create all conda environments!")
+        if names:
+            self.success(f"Created environment{'s' if len(names) > 1 else ''} {' '.join(names)}!")
+        else:
+            self.success("All conda environments created!")
+
+    def activate(
+        self,
+        name: str = typer.Argument(..., help="The name of the environment."),
+        verbose: bool = typer.Option(
+            False, "--verbose", "-v", help="Print the activation command."
+        ),
+    ):
+        env_path = self._get_conda_env_path(name)
+        self.log(f"Activating {name} environment... (type 'exit' to deactivate)")
+        env = Env(self.snakemake_workflow, env_file=env_path.resolve())
+        env.create()
+        user_shell = os.environ.get("SHELL", "/bin/bash")
+        activate_cmd = self._shellcmd(env.address, user_shell)
+        if verbose:
+            self.log(activate_cmd)
+        subprocess.run(activate_cmd, shell=True, env=os.environ.copy(), executable=user_shell)
+        self.log(f"Exiting {name} environment...")