Support global --profile option in modal CLI (#53843)

GitOrigin-RevId: d8556dd8e521bfc10844eee02fb95cbd40bc4f05

Author: mwaskom

py/CHANGELOG_DEV.md

@@ -4,3 +4,4 @@
 
 - Added `Sandbox.logs`(/docs/sdk/py/latest/Sandbox#logs) namespace to retrieve Sandbox entrypoint logs directly from the SDK. The namespace has two different methods, allowing you `fetch()` logs from a specific date/time range, or `tail()` the most recent logs.
 - Added support for setting the default member Role when creating Restricted Environments through the Python SDK and CLI.
+- The `modal` CLI now accepts a global `--profile` option for simpler ad hoc profile selection.

py/modal/cli/_help.py

@@ -7,6 +7,7 @@
 import os
 import shutil
 import sys
+from collections.abc import Sequence
 from typing import Any
 
 import click
@@ -28,6 +29,83 @@
 _HELP_PADDING = 1
 
 
+class ModalGlobalOption(click.Option):
+    """An option supported at every level of the Modal CLI."""
+
+    def __init__(self, *args: Any, environment_variable: str, **kwargs: Any) -> None:
+        kwargs.setdefault("expose_value", False)
+        super().__init__(*args, **kwargs)
+        self.environment_variable = environment_variable
+
+    def set_environment_value(self, value: object) -> None:
+        os.environ[self.environment_variable] = str(value)
+
+
+class ModalProfileOption(ModalGlobalOption):
+    """Global option that updates Modal's active-profile cache."""
+
+    def set_environment_value(self, value: object) -> None:
+        from modal.config import _set_profile
+
+        # The profile is a global variable and currently set when modal is imported
+        # We could probably clean this up and simplify to just setting MODAL_PROFILE
+        _set_profile(str(value))
+
+
+def _root_global_options(ctx: click.Context) -> list[tuple[ModalGlobalOption, click.Context]]:
+    root_ctx = ctx.find_root()
+    return [
+        (param, root_ctx) for param in root_ctx.command.get_params(root_ctx) if isinstance(param, ModalGlobalOption)
+    ]
+
+
+def _global_option_token_length(ctx: click.Context, args: list[str], index: int) -> int:
+    option_name = args[index].split("=", 1)[0]
+    for option, _ in _root_global_options(ctx):
+        if option_name in (*option.opts, *option.secondary_opts):
+            return 1 if "=" in args[index] or option.is_flag else option.nargs + 1
+    return 0
+
+
+def _consume_global_options(ctx: click.Context, args: list[str]) -> list[str]:
+    remaining_args: list[str] = []
+    value: str | bool
+    index = 0
+    while index < len(args):
+        arg = args[index]
+        option_name, separator, option_value = arg.partition("=")
+        global_option = next(
+            (
+                option
+                for option, _ in _root_global_options(ctx)
+                if option_name in (*option.opts, *option.secondary_opts)
+            ),
+            None,
+        )
+        if global_option is None:
+            remaining_args.append(arg)
+            index += 1
+            continue
+
+        if separator:
+            value = option_value
+            index += 1
+        elif global_option.is_flag:
+            value = global_option.flag_value
+            if option_name in global_option.secondary_opts and isinstance(value, bool):
+                value = not value
+            index += 1
+        else:
+            if index + global_option.nargs >= len(args):
+                raise click.UsageError(f"Option '{option_name}' requires an argument.", ctx)
+            value = args[index + 1]
+            index += global_option.nargs + 1
+
+        global_option.set_environment_value(value)
+
+    return remaining_args
+
+
 def use_rich_style() -> bool:
     """Whether help output should be rendered in the rich style."""
     env = os.environ.get("MODAL_RICH_CLI")  # TODO move to config
@@ -82,11 +160,11 @@ def _option_label(param: click.Parameter, ctx: click.Context) -> Text:
 
 def _build_options(cmd: click.Command, ctx: click.Context) -> RenderableType | None:
     rows: list[tuple[Text, str]] = []
-    for param in cmd.get_params(ctx):
-        rec = param.get_help_record(ctx)
+    for param, param_ctx in _options_with_global_options(cmd, ctx):
+        rec = param.get_help_record(param_ctx)
         if rec is None:  # skips arguments and hidden options
             continue
-        rows.append((_option_label(param, ctx), rec[1] or ""))
+        rows.append((_option_label(param, param_ctx), rec[1] or ""))
     if not rows:
         return None
 
@@ -98,6 +176,36 @@ def _build_options(cmd: click.Command, ctx: click.Context) -> RenderableType | N
     return Group(Text("Options", style=_HEADING_STYLE), table)
 
 
+def _global_options(ctx: click.Context) -> Sequence[tuple[click.Option, click.Context]]:
+    root_ctx = ctx.find_root()
+    if root_ctx is ctx:
+        return []
+    return _root_global_options(ctx)
+
+
+def _options_with_global_options(cmd: click.Command, ctx: click.Context) -> list[tuple[click.Parameter, click.Context]]:
+    params = [(param, ctx) for param in cmd.get_params(ctx)]
+    if global_options := _global_options(ctx):
+        for index, (param, _) in enumerate(params):
+            if "--help" in param.opts or "--help" in param.secondary_opts:
+                params[index:index] = global_options
+                break
+        else:
+            params.extend(global_options)
+    return params
+
+
+def _format_options(cmd: click.Command, ctx: click.Context, formatter: click.HelpFormatter) -> None:
+    records = [
+        rec
+        for param, param_ctx in _options_with_global_options(cmd, ctx)
+        if (rec := param.get_help_record(param_ctx)) is not None
+    ]
+    if records:
+        with formatter.section("Options"):
+            formatter.write_dl(records)
+
+
 def _build_epilog(cmd: click.Command) -> RenderableType | None:
     if not cmd.epilog:
         return None
@@ -114,6 +222,10 @@ def group_commands_by_panel(group: click.Group) -> dict[str, list[tuple[str, cli
     return panels
 
 
+def _has_visible_commands(group: click.Group) -> bool:
+    return bool(group_commands_by_panel(group))
+
+
 def _build_commands(group: click.Group, available_width: int) -> RenderableType | None:
     panels = group_commands_by_panel(group)
     if not panels:
@@ -189,6 +301,9 @@ def __init__(self, *args: Any, panel: str | None = None, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
         self.panel = panel
 
+    def format_options(self, ctx: click.Context, formatter: click.HelpFormatter) -> None:
+        _format_options(self, ctx, formatter)
+
     def format_help(self, ctx: click.Context, formatter: click.HelpFormatter) -> None:
         if not use_rich_style():
             return super().format_help(ctx, formatter)
@@ -210,6 +325,7 @@ class ModalGroup(click.Group):
 
     command_class = ModalCommand
     group_class = type  # nested @group.group() reuses the enclosing class
+    defer_global_option_parsing = False
 
     def __init__(self, *args: Any, panel: str | None = None, **kwargs: Any) -> None:
         # Default to showing help when a group is invoked with no subcommand.
@@ -232,6 +348,29 @@ def add_command(
         if hidden is not None:
             cmd.hidden = hidden
 
+    def parse_args(self, ctx: click.Context, args: list[str]) -> list[str]:
+        index = 0
+        while index < len(args):
+            arg = args[index]
+            if global_option_length := _global_option_token_length(ctx, args, index):
+                index += global_option_length
+                continue
+            command = self.commands.get(arg)
+            if command and getattr(command, "defer_global_option_parsing", False):
+                args[:] = _consume_global_options(ctx, args[:index]) + args[index:]
+                break
+            index += 1
+        else:
+            args[:] = _consume_global_options(ctx, args)
+
+        return super().parse_args(ctx, args)
+
+    def format_options(self, ctx: click.Context, formatter: click.HelpFormatter) -> None:
+        if _has_visible_commands(self):
+            self.format_commands(ctx, formatter)
+        else:
+            _format_options(self, ctx, formatter)
+
     def format_commands(self, ctx: click.Context, formatter: click.HelpFormatter) -> None:
         # Replaces click's single flat "Commands:" section with one section per
         # panel so the simple-style help output still preserves grouping.
@@ -249,7 +388,7 @@ def format_help(self, ctx: click.Context, formatter: click.HelpFormatter) -> Non
             [
                 _build_usage(self, ctx),
                 _build_help_text(self),
-                _build_options(self, ctx),
+                None if _has_visible_commands(self) else _build_options(self, ctx),
                 _build_commands(self, _available_width(console)),
                 _build_epilog(self),
             ],

py/modal/cli/config.py

@@ -3,7 +3,8 @@
 
 import click
 
-from modal.config import _profile, _store_user_config, config
+from modal import config as config_module
+from modal.config import _store_user_config, config
 from modal.environments import Environment
 from modal.output import OutputManager
 
@@ -46,7 +47,7 @@ def set_environment(environment_name: str):
     # Confirm that the environment exists by looking it up
     Environment.from_name(environment_name).hydrate()
     _store_user_config({"environment": environment_name})
-    click.echo(f"New default environment for profile {_profile}: {environment_name}")
+    click.echo(f"New default environment for profile {config_module._profile}: {environment_name}")
 
 
 @config_cli.command("set", hidden=True, no_args_is_help=True)

py/modal/cli/entry_point.py

@@ -8,7 +8,7 @@
 from modal.output import OutputManager
 
 from . import run, shell as shell_module
-from ._help import ModalCommand, ModalGroup
+from ._help import ModalCommand, ModalGroup, ModalProfileOption
 from .app import app_cli
 from .billing import billing_cli
 from .bootstrap import bootstrap
@@ -52,6 +52,12 @@
     expose_value=False,
     callback=lambda ctx, param, value: _version_callback(ctx, value),
 )
+@click.option(
+    "--profile",
+    cls=ModalProfileOption,
+    environment_variable="MODAL_PROFILE",
+    help="Use this Modal profile for the command.",
+)
 def entrypoint_cli():
     pass
 
@@ -88,14 +94,13 @@ def check_path():
 
 
 @click.command("setup", cls=ModalCommand, help="Bootstrap Modal's configuration.")
-@click.option("--profile", default=None)
 @synchronizer.create_blocking
-async def setup(profile: str | None = None):
+async def setup():
     check_path()
     print_logo()
 
     # Fetch a new token (same as `modal token new` but redirect to /home once finishes)
-    await _new_token(profile=profile, next_url="/home")
+    await _new_token(next_url="/home")
 
     output = OutputManager.get()
     output.print("[green]→[/green] Run [bold]modal skills install[/bold] to install agent skills")

py/modal/cli/profile.py

@@ -8,11 +8,11 @@
 from rich.json import JSON
 from rich.table import Table
 
+from modal import config as config_module
 from modal._utils.async_utils import synchronizer
 from modal.config import (
     Config,
     _lookup_workspace,
-    _profile,
     config_profiles,
     config_set_active_profile,
 )
@@ -41,7 +41,7 @@ def activate(profile: str):
 
 @profile_cli.command("current", help="Print the currently active Modal profile.")
 def current():
-    click.echo(_profile)
+    click.echo(config_module._profile)
 
 
 @profile_cli.command("list", help="Show all Modal profiles and highlight the active one.")
@@ -62,7 +62,7 @@ async def list_(json: bool | None = False):
 
     rows = []
     for profile, resp in zip(profiles, responses):
-        active = profile == _profile
+        active = profile == config_module._profile
         if isinstance(resp, AuthError):
             workspace = "Unknown (authentication failure)"
         elif isinstance(resp, TimeoutError):
@@ -80,7 +80,7 @@ async def list_(json: bool | None = False):
     if "MODAL_TOKEN_ID" in os.environ:
         try:
             env_based_resp = await _lookup_workspace(
-                config.get("server_url", profile=_profile),
+                config.get("server_url", profile=config_module._profile),
                 os.environ["MODAL_TOKEN_ID"],
                 os.environ.get("MODAL_TOKEN_SECRET", ""),
             )

py/modal/cli/run.py

@@ -23,7 +23,7 @@
 from ..output import OutputManager
 from ..runner import DEPLOYMENT_STRATEGY_TYPE, deploy_app, run_app
 from ..serving import serve_app
-from ._help import ModalCommand, ModalGroup
+from ._help import ModalCommand, ModalGroup, _consume_global_options, _global_option_token_length
 from .import_refs import (
     CLICommand,
     MethodReference,
@@ -411,6 +411,33 @@ def _get_runnable_list(all_usable_commands: list[CLICommand]) -> str:
 class RunGroup(ModalGroup):
     """Click group that resolves subcommands dynamically from a file/module ref."""
 
+    defer_global_option_parsing = True
+
+    def parse_args(self, ctx: click.Context, args: list[str]) -> list[str]:
+        option_params: dict[str, click.Option] = {}
+        for param in self.get_params(ctx):
+            if isinstance(param, click.Option):
+                for option in (*param.opts, *param.secondary_opts):
+                    option_params[option] = param
+
+        func_ref_index = 0
+        while func_ref_index < len(args):
+            arg = args[func_ref_index]
+            if global_option_length := _global_option_token_length(ctx, args, func_ref_index):
+                func_ref_index += global_option_length
+            elif arg == "--":
+                func_ref_index += 1
+                break
+            elif param := option_params.get(arg.split("=", 1)[0]):
+                func_ref_index += 1 if "=" in arg or param.is_flag else param.nargs + 1
+            elif arg.startswith("-"):
+                func_ref_index += 1
+            else:
+                break
+
+        args[:] = _consume_global_options(ctx, args[:func_ref_index]) + args[func_ref_index:]
+        return click.Group.parse_args(self, ctx, args)
+
     def get_command(self, ctx: click.Context, cmd_name: str) -> click.Command | None:
         # note: get_command here is run before the "group logic" in the `run` logic below
         # so to ensure that `env` has been globally populated before user code is loaded, it