➡️➡️ Tab Tab for zsh: CLI completion generator (#208)

Co-authored-by: Anthony Wu <pls-file-gh-issue@users.noreply.github.com>
Co-authored-by: filipstrand <strand.filip@gmail.com>

README.md

@@ -12,6 +12,7 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black
 
 - [Philosophy](#philosophy)
 - [💿 Installation](#-installation)
+- [🚀 Shell Completions (Quick Start)](#-shell-completions-quick-start)
 - [🖼️ Generating an image](#%EF%B8%8F-generating-an-image)
   * [📜 Full list of Command-Line Arguments](#-full-list-of-command-line-arguments)
 - [⏱️ Image generation speed (updated)](#%EF%B8%8F-image-generation-speed-updated)
@@ -139,6 +140,31 @@ pip install -U mflux
 
 *If you have trouble installing MFLUX, please see the [installation related issues section](https://github.com/filipstrand/mflux/issues?q=is%3Aissue+install+).* 
 
+### ⌨️ Shell Completions (Quick Start)
+
+MFLUX supports ZSH (default on macOS) shell completions for all CLI commands.
+
+To enable completions:
+
+```sh
+# Install completions
+mflux-completions
+
+# Reload your shell
+exec zsh
+
+# Now try tab completion!
+mflux-generate --<TAB>
+```
+
+If you encounter issues (ZSH completion setup can vary by system), check your installation:
+
+```sh
+mflux-completions --check
+```
+
+For more details and troubleshooting, see the [completions documentation](src/mflux/completions/README.md).
+
 ### 🖼️ Generating an image
 
 Run the command `mflux-generate` by specifying a prompt and the model and some optional arguments. For example, here we use a quantized version of the `schnell` model for 2 steps:

pyproject.toml

@@ -71,6 +71,7 @@ mflux-save-depth = "mflux.save_depth:main"
 mflux-train = "mflux.train:main"
 mflux-upscale = "mflux.upscale:main"
 mflux-lora-library = "mflux.lora_library:main"
+mflux-completions = "mflux.completions.install:main"
 
 [tool.setuptools.packages.find]
 where = ["src"]

src/mflux/completions/README.md

@@ -0,0 +1,144 @@
+# mflux Shell Completions
+
+This module provides ZSH shell completion support for all mflux CLI commands.
+
+## Quick Start
+
+1. **Install completions:**
+   ```bash
+   mflux-completions
+   ```
+
+2. **Reload your shell:**
+   ```bash
+   exec zsh
+   ```
+
+3. **Test it:**
+   ```bash
+   mflux-generate --<TAB>
+   ```
+
+## Features
+
+- **Auto-completion for all 15 mflux commands**: Generate, train, save, upscale, and more
+- **Smart completions**: 
+  - Model names (dev, schnell, dev-fill, or HuggingFace repos)
+  - Quantization levels (3, 4, 6, 8)
+  - LoRA styles (storyboard, portrait, etc.)
+  - File paths with native completion
+  - Common percentage values for battery limit
+- **Dynamic generation**: Completions stay in sync with code changes
+- **Mutually exclusive options**: Properly handles -m/--model style options
+
+## Installation
+
+### Automatic Installation
+
+The simplest way is to let mflux-completions auto-detect the best location:
+
+```bash
+mflux-completions
+```
+
+### Custom Directory
+
+Install to a specific directory:
+
+```bash
+mflux-completions --dir ~/.config/zsh/completions
+```
+
+### Update Existing Installation
+
+Update an existing completion file:
+
+```bash
+mflux-completions --update
+```
+
+## Troubleshooting
+
+### Check Installation
+
+Verify that completions are properly installed:
+
+```bash
+mflux-completions --check
+```
+
+This will:
+- Check if the completion file exists in your $fpath
+- Verify that compinit is available
+- Provide troubleshooting steps if needed
+
+### Common Issues
+
+1. **Completions not working after installation:**
+   - Start a new shell: `exec zsh`
+
+2. **Completions not updated after mflux upgrade:**
+   - Update completions: `mflux-completions --update`
+   - Then force reload: `compinit -u`
+
+3. **Permission denied:**
+   - Use `--dir` to specify a user-writable directory
+
+4. **Completion shows \} or other artifacts:**
+   - Update to the latest version: `mflux-completions --update`
+
+For Zsh-specific configuration issues (fpath, compinit, etc.), please refer to the [official Zsh documentation](https://zsh.sourceforge.io/Doc/Release/Completion-System.html).
+
+### Manual Installation
+
+If automatic installation fails:
+
+```bash
+# Generate completion script
+mflux-completions --print > _mflux
+
+# Move to a directory in your $fpath
+mkdir -p ~/.zsh/completions
+mv _mflux ~/.zsh/completions/
+
+# Restart your shell
+exec zsh
+```
+
+## How It Works
+
+The completion system:
+
+1. **Introspects argparse**: Dynamically extracts all command-line options from the parsers
+2. **Generates ZSH functions**: Creates completion functions for each mflux command
+3. **Provides context-aware completions**: Different completion strategies based on argument type
+4. **Handles special cases**: Custom completions for models, LoRA styles, etc.
+
+## Development
+
+### Adding New Commands
+
+When adding a new mflux CLI command:
+
+1. Create the command with argparse as usual
+2. Add it to the `commands` list in `generator.py`
+3. Add the parser creation logic to `create_parser_for_command()`
+4. Completions will be automatically generated!
+
+### Testing Completions
+
+Test the generator without installing:
+
+```bash
+python -m mflux.completions.generator
+```
+
+### File Structure
+
+- `generator.py`: Core logic for generating ZSH completion scripts
+- `install.py`: Installation utility with user-friendly CLI
+- `README.md`: This documentation
+
+## Platform Support
+
+Currently supports ZSH (the default macOS shell since 2019). Supporting other shells is possible but is considered out of scope for the official repository.

src/mflux/completions/__init__.py

@@ -0,0 +1 @@
+"""ZSH completion generator for mflux CLI commands."""

src/mflux/completions/generator.py

@@ -0,0 +1,328 @@
+"""Generate ZSH completion scripts for mflux commands."""
+
+import argparse
+from pathlib import Path
+
+from mflux.community.in_context.utils.in_context_loras import LORA_NAME_MAP
+from mflux.ui import defaults as ui_defaults
+from mflux.ui.cli.parsers import CommandLineParser
+
+
+class CompletionGenerator:
+    """Generate ZSH completion scripts by introspecting argparse parsers."""
+
+    def __init__(self):
+        self.commands = [
+            "mflux-generate",
+            "mflux-generate-controlnet",
+            "mflux-generate-in-context",
+            "mflux-generate-in-context-edit",
+            "mflux-generate-in-context-catvton",
+            "mflux-generate-fill",
+            "mflux-generate-depth",
+            "mflux-generate-redux",
+            "mflux-concept",
+            "mflux-concept-from-image",
+            "mflux-save",
+            "mflux-save-depth",
+            "mflux-train",
+            "mflux-upscale",
+            "mflux-lora-library",
+        ]
+
+    def create_parser_for_command(self, command: str) -> CommandLineParser:
+        """Create the appropriate parser for a given command."""
+        parser = CommandLineParser(prog=command, add_help=False)
+
+        if command == "mflux-generate":
+            parser.add_general_arguments()
+            parser.add_model_arguments()
+            parser.add_lora_arguments()
+            parser.add_image_generator_arguments(supports_metadata_config=True)
+            parser.add_image_to_image_arguments()
+            parser.add_image_outpaint_arguments()
+            parser.add_output_arguments()
+
+        elif command == "mflux-generate-controlnet":
+            parser.add_general_arguments()
+            parser.add_model_arguments()
+            parser.add_lora_arguments()
+            parser.add_controlnet_arguments()
+            parser.add_image_generator_arguments()
+            parser.add_output_arguments()
+
+        elif command == "mflux-generate-in-context":
+            parser.add_general_arguments()
+            parser.add_model_arguments()
+            parser.add_image_generator_arguments()
+            parser.add_image_to_image_arguments(required=True)
+            parser.add_output_arguments()
+            # Add save-full-image manually since it's not in parsers.py
+            parser.add_argument("--save-full-image", action="store_true")
+
+        elif command == "mflux-generate-in-context-edit":
+            parser.add_general_arguments()
+            parser.add_model_arguments(require_model_arg=False)
+            parser.add_lora_arguments()
+            parser.add_image_generator_arguments(supports_metadata_config=False, require_prompt=False)
+            parser.add_in_context_edit_arguments()
+            parser.add_in_context_arguments()
+            parser.add_output_arguments()
+
+        elif command == "mflux-generate-in-context-catvton":
+            parser.add_general_arguments()
+            parser.add_model_arguments(require_model_arg=False)
+            parser.add_lora_arguments()
+            parser.add_image_generator_arguments(supports_metadata_config=False, require_prompt=False)
+            parser.add_catvton_arguments()
+            parser.add_in_context_arguments()
+            parser.add_output_arguments()
+
+        elif command == "mflux-generate-fill":
+            parser.add_general_arguments()
+            parser.add_model_arguments()
+            parser.add_lora_arguments()
+            parser.add_image_generator_arguments()
+            parser.add_fill_arguments()
+            parser.add_output_arguments()
+
+        elif command == "mflux-generate-depth":
+            parser.add_general_arguments()
+            parser.add_model_arguments()
+            parser.add_lora_arguments()
+            parser.add_image_generator_arguments()
+            parser.add_depth_arguments()
+            parser.add_output_arguments()
+
+        elif command == "mflux-generate-redux":
+            parser.add_general_arguments()
+            parser.add_model_arguments()
+            parser.add_lora_arguments()
+            parser.add_image_generator_arguments()
+            parser.add_redux_arguments()
+            parser.add_output_arguments()
+
+        elif command == "mflux-concept":
+            parser.add_general_arguments()
+            parser.add_model_arguments()
+            parser.add_concept_attention_arguments()
+            parser.add_image_generator_arguments()
+            parser.add_output_arguments()
+
+        elif command == "mflux-concept-from-image":
+            parser.add_general_arguments()
+            parser.add_model_arguments()
+            parser.add_concept_from_image_arguments()
+            parser.add_image_generator_arguments()
+            parser.add_output_arguments()
+
+        elif command == "mflux-save":
+            parser.add_model_arguments(path_type="save")
+
+        elif command == "mflux-save-depth":
+            parser.add_save_depth_arguments()
+            parser.add_output_arguments()
+
+        elif command == "mflux-train":
+            parser.add_model_arguments(require_model_arg=False)
+            parser.add_training_arguments()
+
+        elif command == "mflux-upscale":
+            parser.add_general_arguments()
+            parser.add_model_arguments()
+            parser.add_image_generator_arguments(supports_metadata_config=True)
+            parser.add_image_to_image_arguments(required=True)
+            parser.add_output_arguments()
+
+        elif command == "mflux-lora-library":
+            # Special case: has subcommands
+            subparsers = parser.add_subparsers(dest="subcommand")
+            list_parser = subparsers.add_parser("list")
+            list_parser.add_argument("--paths", action="store_true")
+
+        return parser
+
+    def escape_description(self, desc: str) -> str:
+        """Escape special characters in descriptions for ZSH."""
+        if not desc:
+            return ""
+        # Escape brackets, quotes, and other special characters
+        desc = desc.replace("\\", "\\\\")  # Escape backslashes first
+        desc = desc.replace("[", "\\[").replace("]", "\\]")
+        desc = desc.replace("'", "")  # Remove single quotes entirely
+        desc = desc.replace('"', '\\"')
+        desc = desc.replace("`", "\\`")
+        desc = desc.replace("$", "\\$")
+        return desc
+
+    def format_argument_spec(self, action: argparse.Action) -> list[str]:
+        """Format an argparse action into ZSH completion syntax."""
+        specs = []
+
+        # Handle options with both short and long forms
+        if action.option_strings:
+            opts = action.option_strings
+            desc = self.escape_description(action.help or "")
+
+            # For store_true/store_false actions, no value spec needed
+            is_flag = isinstance(action, (argparse._StoreTrueAction, argparse._StoreFalseAction))
+
+            if len(opts) == 2:
+                # Both short and long form - create exclusion group
+                opt_spec = f"'({opts[0]} {opts[1]})'{{{opts[0]},{opts[1]}}}"
+            else:
+                # Only one form
+                opt_spec = f"'{opts[0]}'"
+
+            opt_spec += f"'[{desc}]"
+
+            # Add value spec for non-flag arguments
+            if not is_flag and (action.type or action.choices or action.nargs):
+                opt_spec += self.get_value_spec(action)
+
+            opt_spec += "'"
+            specs.append(opt_spec)
+
+        return specs
+
+    def get_value_spec(self, action: argparse.Action) -> str:
+        """Get the value specification for an argument."""
+        # Check for special cases first
+        if action.dest == "model":
+            return ":model:_mflux_models"
+        elif action.dest == "quantize":
+            return ":quantization:_mflux_quantize"
+        elif action.dest == "lora_style":
+            return ":style:_mflux_lora_styles"
+
+        if action.choices:
+            # Fixed choices
+            choices = " ".join(str(c) for c in action.choices)
+            return f":value:({choices})"
+
+        elif action.type == Path or (action.dest and ("path" in action.dest or "file" in action.dest)):
+            # File completion
+            return ":file:_files"
+
+        elif action.type is int:
+            # Integer value
+            return ":number:"
+
+        elif action.type is float:
+            # Float value
+            return ":float:"
+
+        elif callable(action.type):
+            # Callable type (like lambda) - usually returns int or float
+            # Check the dest name for hints
+            if "percentage" in (action.dest or "") or "percent" in (action.dest or ""):
+                return ":percentage:(1 5 10 15 20 25 30 40 50 60 70 80 90 95 99)"
+            return ":value:"
+
+        elif action.nargs in ["*", "+"]:
+            # Multiple values
+            if action.type == Path or (action.dest and ("path" in action.dest)):
+                return ":files:_files"
+            return ":values:"
+
+        elif action.type is str or action.type is None:
+            # String value
+            return f":{action.dest or 'value'}:"
+
+        return ""
+
+    def generate_command_function(self, command: str, parser: CommandLineParser) -> str:
+        """Generate the ZSH completion function for a specific command."""
+        func_name = command.replace("-", "_")
+        lines = [f"_{func_name}() {{"]
+        lines.append("    local -a args")
+        lines.append("    args=(")
+
+        # Extract all actions from the parser
+        for action in parser._actions:
+            if isinstance(action, argparse._HelpAction):
+                continue
+            arg_specs = self.format_argument_spec(action)
+            lines.extend(f"        {spec}" for spec in arg_specs)
+
+        lines.append("    )")
+        lines.append("    _arguments -s -S $args")
+        lines.append("}")
+        lines.append("")
+
+        return "\n".join(lines)
+
+    def generate_header(self) -> str:
+        """Generate the completion script header."""
+        commands = " ".join(self.commands)
+        return f"""#compdef {commands}
+# ZSH completion for mflux commands
+# Generated by mflux-completions
+
+"""
+
+    def generate_helper_functions(self) -> str:
+        """Generate helper functions for common completions."""
+        helpers = []
+
+        # Model completion helper
+        model_choices = " ".join(f"'{m}[{m} model]'" for m in ui_defaults.MODEL_CHOICES)
+        helpers.append(f"""_mflux_models() {{
+    _values 'model' \\
+        {model_choices} \\
+        '*:Hugging Face repo:(org/model)'
+}}
+""")
+
+        # LoRA style completion helper
+        if LORA_NAME_MAP:
+            lora_choices = " ".join(f"'{k}[{v}]'" for k, v in LORA_NAME_MAP.items())
+            helpers.append(f"""_mflux_lora_styles() {{
+    _values 'lora style' \\
+        {lora_choices}
+}}
+""")
+
+        # Quantization choices
+        quant_choices = " ".join(f"'{q}[{q}-bit quantization]'" for q in ui_defaults.QUANTIZE_CHOICES)
+        helpers.append(f"""_mflux_quantize() {{
+    _values 'quantization' \\
+        {quant_choices}
+}}
+""")
+
+        return "\n".join(helpers)
+
+    def generate_main_function(self) -> str:
+        """Generate the main completion dispatcher."""
+        lines = ["# Main completion dispatcher", "_mflux() {", "    local cmd=$words[1]", "    case $cmd in"]
+
+        for command in self.commands:
+            func_name = command.replace("-", "_")
+            lines.append(f"        {command})")
+            lines.append(f"            _{func_name}")
+            lines.append("            ;;")
+
+        lines.extend(["        *)", "            _default", "            ;;", "    esac", "}", "", "_mflux"])
+
+        return "\n".join(lines)
+
+    def generate(self) -> str:
+        """Generate the complete ZSH completion script."""
+        script = [self.generate_header()]
+        script.append(self.generate_helper_functions())
+
+        # Generate completion functions for each command
+        for command in self.commands:
+            parser = self.create_parser_for_command(command)
+            script.append(self.generate_command_function(command, parser))
+
+        script.append(self.generate_main_function())
+
+        return "\n".join(script)
+
+
+if __name__ == "__main__":
+    # Test the generator by printing out the completion function to stdout
+    generator = CompletionGenerator()
+    print(generator.generate())

src/mflux/completions/install.py

@@ -0,0 +1,203 @@
+#!/usr/bin/env python3
+"""Install ZSH completions for mflux commands."""
+
+import argparse
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+from mflux.completions.generator import CompletionGenerator
+
+
+def get_zsh_fpath():
+    """Get the ZSH fpath directories."""
+    try:
+        result = subprocess.run(
+            ["zsh", "-c", "echo $fpath"],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+        return result.stdout.strip().split()
+    except subprocess.CalledProcessError:
+        return []
+
+
+def find_completion_dir():
+    """Find appropriate directory for completion files."""
+    # Common completion directories in order of preference
+    candidates = [
+        Path.home() / ".zsh" / "completions",
+        Path.home() / ".config" / "zsh" / "completions",
+        Path("/usr/local/share/zsh/site-functions"),
+        Path("/opt/homebrew/share/zsh/site-functions"),
+    ]
+
+    # Check fpath directories
+    fpath = get_zsh_fpath()
+    for path_str in fpath:
+        path = Path(path_str)
+        if path.exists() and os.access(path, os.W_OK):
+            candidates.insert(0, path)
+
+    # Return first writable directory
+    for candidate in candidates:
+        if candidate.exists() and os.access(candidate, os.W_OK):
+            return candidate
+
+    # Default to user directory
+    user_dir = Path.home() / ".zsh" / "completions"
+    user_dir.mkdir(parents=True, exist_ok=True)
+    return user_dir
+
+
+def check_installation():
+    """Check if completions are properly installed and accessible."""
+    print("Checking mflux completions installation...\n")
+
+    # Check if completion file exists
+    found_locations = []
+    fpath = get_zsh_fpath()
+
+    for path_str in fpath:
+        path = Path(path_str)
+        completion_file = path / "_mflux"
+        if completion_file.exists():
+            found_locations.append(completion_file)
+
+    if not found_locations:
+        print("❌ No mflux completion file found in $fpath")
+        print("\nSearched in:")
+        for path_str in fpath[:5]:  # Show first 5 locations
+            print(f"  - {path_str}")
+        print("\nRun 'mflux-completions' to install completions")
+        return False
+
+    print("✓ Found completion file(s):")
+    for location in found_locations:
+        print(f"  - {location}")
+
+    # Check if compinit is loaded in interactive shell
+    try:
+        result = subprocess.run(
+            ["zsh", "-i", "-c", "type compinit"],
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+        if result.returncode == 0:
+            print("\n✓ compinit is available")
+        else:
+            print("\n⚠️  compinit may not be initialized")
+            print("   If completions don't work, add to ~/.zshrc:")
+            print("   autoload -U compinit && compinit")
+    except (subprocess.SubprocessError, OSError):
+        print("\n⚠️  Could not verify compinit")
+
+    # Test a completion
+    print("\n✓ Installation appears correct!")
+    print("\nTo test completions:")
+    print("  1. Start a new shell or run: exec zsh")
+    print("  2. Type: mflux-generate --<TAB>")
+    print("\nIf completions don't work, try:")
+    print("  rm -f ~/.zcompdump && compinit")
+
+    return True
+
+
+def main():
+    """Main entry point for completion installation."""
+    parser = argparse.ArgumentParser(
+        description="Install or generate ZSH completions for mflux commands",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Install to default location
+  mflux-completions
+
+  # Generate to stdout
+  mflux-completions --print
+
+  # Install to specific directory
+  mflux-completions --dir ~/.config/zsh/completions
+
+  # Update existing completion
+  mflux-completions --update
+
+  # Check if completions are properly installed
+  mflux-completions --check
+""",
+    )
+
+    parser.add_argument(
+        "--print",
+        action="store_true",
+        help="Print completion script to stdout instead of installing",
+    )
+    parser.add_argument(
+        "--dir",
+        type=Path,
+        help="Directory to install completion file (default: auto-detect)",
+    )
+    parser.add_argument(
+        "--update",
+        action="store_true",
+        help="Update existing completion file",
+    )
+    parser.add_argument(
+        "--check",
+        action="store_true",
+        help="Check if completions are properly installed",
+    )
+
+    args = parser.parse_args()
+
+    # Handle --check flag
+    if args.check:
+        sys.exit(0 if check_installation() else 1)
+
+    # Generate completion script
+    generator = CompletionGenerator()
+    completion_script = generator.generate()
+
+    if args.print:
+        # Just print to stdout
+        print(completion_script)
+        return
+
+    # Determine installation directory
+    if args.dir:
+        install_dir = args.dir
+        install_dir.mkdir(parents=True, exist_ok=True)
+    else:
+        install_dir = find_completion_dir()
+
+    # Install completion file
+    completion_file = install_dir / "_mflux"
+
+    if completion_file.exists() and not args.update:
+        print(f"Error: Completion file already exists at {completion_file}")
+        print("Use --update to overwrite")
+        sys.exit(1)
+
+    try:
+        completion_file.write_text(completion_script)
+        print(f"✓ Installed ZSH completions to: {completion_file}")
+        print()
+        print("To activate completions:")
+        print("1. Ensure this directory is in your $fpath:")
+        print(f"   echo 'fpath=({install_dir} $fpath)' >> ~/.zshrc")
+        print("2. Reload your shell:")
+        print("   exec zsh")
+        print()
+        print("If completions don't work immediately, you may need to:")
+        print("   rm -f ~/.zcompdump && compinit")
+
+    except OSError as e:
+        print(f"Error installing completion file: {e}")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()