Writing Custom Operations

from pyinfra import host
from pyinfra.api import operation, StringCommand

@operation()
def ensure_package(package_name: str):
    """
    Ensure a package is installed.
    
    + package_name: name of the package to install
    """
    # Check current state using a fact
    installed_packages = host.get_fact(InstalledPackages)
    
    # Only yield commands if change is needed
    if package_name not in installed_packages:
        yield StringCommand("apt-get", "install", "-y", package_name)
    else:
        host.noop(f"package {package_name} is already installed")

@operation(
    is_idempotent: bool = True,
    idempotent_notice: Optional[str] = None,
    is_deprecated: bool = False,
    deprecated_for: Optional[str] = None,
)

from pyinfra.api import operation, StringCommand

@operation(is_idempotent=False)
def reboot(delay: int = 10, interval: int = 1, reboot_timeout: int = 300):
    """
    Reboot the server and wait for reconnection.
    
    + delay: number of seconds to wait before attempting reconnect
    + interval: interval (s) between reconnect attempts
    + reboot_timeout: total time before giving up reconnecting
    """
    yield StringCommand("reboot", _success_exit_codes=[0, -1])

from pyinfra.api import StringCommand, QuoteString

# Simple command
yield StringCommand("systemctl", "restart", "nginx")

# Command with quoted arguments (handles spaces and special chars)
yield StringCommand("echo", QuoteString("Hello, World!"), ">>", "/var/log/app.log")

# Command with custom success codes
yield StringCommand("reboot", _success_exit_codes=[0, -1])

from pyinfra.api import FileUploadCommand
from io import StringIO

# Upload from file path
yield FileUploadCommand(
    src="/local/path/config.yml",
    dest="/etc/app/config.yml"
)

# Upload from IO object
config_content = StringIO("server:\n  port: 8080\n")
yield FileUploadCommand(
    src=config_content,
    dest="/etc/app/config.yml"
)

from pyinfra.api import FileDownloadCommand

yield FileDownloadCommand(
    src="/var/log/app.log",
    dest="./local-logs/app.log"
)

from pyinfra.api import FunctionCommand

def cleanup_temp_files(state, host):
    """Python function executed on the pyinfra controller."""
    print(f"Cleaning up for {host.name}")
    # Can access state and host objects
    
yield FunctionCommand(cleanup_temp_files, (), {})

from pyinfra import host
from pyinfra.api import operation
from pyinfra.facts.files import File, Directory
from pyinfra.facts.server import Hostname, User

@operation()
def configure_app(config_path: str):
    """
    Configure application based on host state.
    """
    # Get simple facts
    hostname = host.get_fact(Hostname)
    current_user = host.get_fact(User)
    
    # Get facts with arguments
    config_exists = host.get_fact(File, path=config_path)
    app_dir_exists = host.get_fact(Directory, path="/opt/app")
    
    if not app_dir_exists:
        yield "mkdir -p /opt/app"
    
    if not config_exists:
        yield f"echo 'hostname={hostname}' > {config_path}"

from pyinfra import host, logger
from pyinfra.api import operation, OperationError, StringCommand, QuoteString
from pyinfra.facts.server import Which
from pyinfra.facts.files import File

@operation()
def download(
    src: str,
    dest: str,
    user: str | None = None,
    group: str | None = None,
    mode: str | None = None,
    force: bool = False,
    sha256sum: str | None = None,
):
    """
    Download files from remote locations using curl or wget.
    
    + src: source URL of the file
    + dest: where to save the file
    + user: user to own the file
    + group: group to own the file
    + mode: permissions of the file
    + force: always download the file, even if it already exists
    + sha256sum: sha256 hash to checksum the downloaded file against
    
    **Example:**
    
    .. code:: python
    
        from pyinfra.operations import files
        files.download(
            name="Download config file",
            src="https://example.com/config.yml",
            dest="/etc/app/config.yml",
            mode="644",
        )
    """
    
    info = host.get_fact(File, path=dest)
    
    # Destination is a directory?
    if info is False:
        raise OperationError(
            f"Destination {dest} already exists and is not a file"
        )
    
    # Determine if we need to download
    download = force or info is None
    
    if sha256sum and info:
        from pyinfra.facts.files import Sha256File
        if sha256sum != host.get_fact(Sha256File, path=dest):
            download = True
    
    if download:
        temp_file = host.get_temp_filename(dest)
        
        # Try curl first, then wget
        if host.get_fact(Which, command="curl"):
            yield StringCommand(
                "curl", "-sSLf", QuoteString(src), "-o", QuoteString(temp_file)
            )
        elif host.get_fact(Which, command="wget"):
            yield StringCommand(
                "wget", "-q", QuoteString(src), "-O", QuoteString(temp_file)
            )
        else:
            raise OperationError("Neither curl nor wget is available")
        
        yield StringCommand("mv", QuoteString(temp_file), QuoteString(dest))
        
        if user or group:
            from pyinfra.operations.util.files import chown
            yield chown(dest, user, group)
        
        if mode:
            from pyinfra.operations.util.files import chmod
            yield chmod(dest, mode)
        
        if sha256sum:
            yield (
                f"(sha256sum {dest} | grep {sha256sum}) || "
                f"(echo 'SHA256 did not match!' && exit 1)"
            )
    else:
        host.noop(f"file {dest} has already been downloaded")

@operation()
def conditional_restart(service: str, config_changed: bool):
    """Restart service only if config changed."""
    if config_changed:
        yield f"systemctl restart {service}"
    else:
        host.noop(f"service {service} config unchanged, no restart needed")

from pyinfra.operations import files, server

@operation()
def deploy_app(version: str):
    """Deploy application with specific version."""
    # Use _inner to call operations within operations
    for command in files.directory._inner(
        path="/opt/app",
        user="app",
        group="app",
        mode="755",
    ):
        yield command
    
    for command in files.download._inner(
        src=f"https://releases.example.com/app-{version}.tar.gz",
        dest="/tmp/app.tar.gz",
    ):
        yield command
    
    yield "tar -xzf /tmp/app.tar.gz -C /opt/app"
    yield "rm /tmp/app.tar.gz"

from pyinfra.api import OperationError, OperationTypeError

@operation()
def install_with_validation(package: str, min_disk_gb: int = 10):
    """Install package with disk space validation."""
    if not isinstance(package, str):
        raise OperationTypeError("package must be a string")
    
    # Check available disk space
    from pyinfra.facts.server import Mounts
    mounts = host.get_fact(Mounts)
    root_mount = mounts.get("/")
    
    if root_mount:
        available_gb = root_mount["available"] / (1024**3)
        if available_gb < min_disk_gb:
            raise OperationError(
                f"Insufficient disk space: {available_gb:.1f}GB available, "
                f"{min_disk_gb}GB required"
            )
    
    yield f"apt-get install -y {package}"

from pyinfra import host, state, context

@operation()
def context_aware_operation():
    """Access host and state information during operation."""
    # Access current host
    logger.info(f"Running on host: {host.name}")
    logger.info(f"Host data: {host.data}")
    
    # Access deployment state
    if state.is_executing:
        logger.info("Currently executing operations")
    
    # Access current operation metadata
    if host.current_op_hash:
        logger.info(f"Operation hash: {host.current_op_hash}")
    
    yield "echo 'Context-aware operation'"

# In your deploy script
from pyinfra import host
from pyinfra.operations import custom

# Test with dry-run mode first
# pyinfra inventory.py deploy.py --dry

result = custom.my_operation(
    name="Test my operation",
    param1="value1",
)

# Check if operation would make changes
if result.will_change:
    print("Operation would make changes")

# After execution, check results
if result.did_succeed():
    print(f"Operation succeeded")
    print(f"Changed: {result.did_change()}")
    print(f"Output: {result.stdout}")