Operation

from pyinfra.api import operation

@operation()
def install_package(package: str):
    """Install a package using apt."""
    yield f"apt-get install -y {package}"

from pyinfra.api import operation

@operation()
def create_user_with_key(username: str, ssh_key: str):
    """Create a user and add their SSH key."""
    # Create user
    yield f"useradd -m {username}"
    
    # Create .ssh directory
    yield f"mkdir -p /home/{username}/.ssh"
    
    # Add SSH key
    yield f'echo "{ssh_key}" > /home/{username}/.ssh/authorized_keys'
    
    # Set permissions
    yield f"chmod 700 /home/{username}/.ssh"
    yield f"chmod 600 /home/{username}/.ssh/authorized_keys"
    yield f"chown -R {username}:{username} /home/{username}/.ssh"

@operation()
def simple_command():
    yield "echo 'Hello World'"

from pyinfra.api import StringCommand, QuoteString

@operation()
def safe_command(user_input: str):
    # Automatically quotes and escapes user input
    yield StringCommand("echo", QuoteString(user_input))

from pyinfra.api import FileUploadCommand

@operation()
def upload_config():
    yield FileUploadCommand(
        src="local/config.txt",
        dest="/etc/app/config.txt",
    )

from pyinfra.api import FunctionCommand

def my_callback():
    print("This runs on the host")

@operation()
def with_callback():
    yield "echo 'Before callback'"
    yield FunctionCommand(my_callback, (), {})
    yield "echo 'After callback'"

from pyinfra.operations import apt

# Returns OperationMeta
op = apt.packages(name="Install nginx", packages=["nginx"])

# Check if operation will make changes
if op.will_change:
    print("This operation will modify the system")

from pyinfra.api import add_op

def add_op(
    state: State,
    op_func: Callable,
    *args,
    **kwargs
) -> dict[Host, OperationMeta]:
    """Add an operation to state.
    
    Args:
        state: The deploy state
        op_func: The operation function
        args/kwargs: Passed to the operation
        
    Returns:
        Dictionary mapping hosts to OperationMeta objects
    """

from pyinfra import Config, Inventory, State
from pyinfra.api import add_op
from pyinfra.operations import server

# Create state
inventory = Inventory((["host1", "host2"], {}))
state = State(inventory, Config())
state.init(inventory, Config())

# Add operation to all hosts
results = add_op(
    state,
    server.shell,
    commands=["echo 'Hello'"],
)

# Check results per host
for host, op_meta in results.items():
    print(f"{host}: will_change={op_meta.will_change}")

# Add operation to specific host
host = inventory.get_host("host1")
add_op(
    state,
    server.shell,
    commands=["echo 'Hello host1'"],
    host=host,
)

from pyinfra import host
from pyinfra.facts.server import LinuxDistribution

@operation()
def install_package_for_distro(package: str):
    distro = host.get_fact(LinuxDistribution)
    
    if distro["name"] == "Ubuntu":
        yield f"apt-get install -y {package}"
    elif distro["name"] == "CentOS":
        yield f"yum install -y {package}"

# operations/custom.py
from pyinfra import host
from pyinfra.api import operation, StringCommand, QuoteString
from pyinfra.facts.files import File

@operation(is_idempotent=True)
def configure_app(config_file: str, key: str, value: str):
    """Set a configuration key in an app config file.
    
    Args:
        config_file: Path to config file
        key: Configuration key
        value: Configuration value
    """
    # Check if config file exists
    file_info = host.get_fact(File, path=config_file)
    
    if not file_info:
        # Create config file
        yield f"touch {config_file}"
    
    # Update configuration
    # Use QuoteString to safely handle special characters
    yield StringCommand(
        "sed", "-i",
        f"s/^{key}=.*/{key}={value}/",
        QuoteString(config_file)
    )
    
    # Restart app if config changed
    yield "systemctl restart myapp"

# Usage
configure_app(
    config_file="/etc/myapp/config.ini",
    key="port",
    value="8080",
    _sudo=True,
)