Using guard in Python

2 minute read

Bringing Swift’s Guard Statement to Python

One of the most convenient constructs in Swift is the guard statement. It provides a clear and readable way to check preconditions before executing code. If a condition isn’t met, the guard statement allows early exit, improving readability and reducing nested code blocks. Unfortunately, Python does not have a native guard statement. However, I decided to create a decorator-based implementation that offers similar functionality.

Implementing a Guard Decorator in Python

Since Python does not support inline guard statements like Swift, the next best approach is using decorators. This method allows us to declare preconditions above a function definition, making it clear which conditions must be satisfied before execution.

from functools import wraps
from typing import Any, Callable

def guard(condition_func, exception_func=lambda: ValueError("Guard condition failed"), raiseExceptionOnException=False):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            if raiseExceptionOnException:
                try:
                    if not condition_func(*args, **kwargs):
                        raise exception_func()
                except Exception:
                    raise exception_func()
            else:
                if not condition_func(*args, **kwargs):
                    raise exception_func()
            return func(*args, **kwargs)
        return wrapper
    return decorator

def guard_input(condition_func: Callable[..., bool], message: str):
    def decorator(func: Callable[..., Any]):
        return guard(condition_func, lambda: ValueError(message))(func)
    return decorator

How It Works

The guard decorator takes:

condition_func: A function that determines whether execution should proceed.
exception_func: A function that returns an exception to be raised if the condition fails (defaulting to ValueError).
raiseExceptionOnException: If True, exceptions raised during condition evaluation trigger the exception function.

This approach makes it explicit which conditions must be met before executing a function.

Demonstrating the Guard Decorator with Tests

The following tests not only validate the guard decorator but also serve as a practical demonstration of how it works in different scenarios using pytest:

import pytest
from system.guard import guard

class CustomGuardTestException(Exception):
    pass

@guard(lambda value: value is not None, lambda: CustomGuardTestException("Value is None"))
@guard(lambda value: value != "TEST", lambda: CustomGuardTestException("Value is Test, this is not allowed"))
def example_function(value):
    print(f"Value is {value}")

# Demonstrating valid and invalid function calls

def test_guard_succeeds():
    """Demonstrates that the function executes normally when conditions are met."""
    example_function("TEST2")
    example_function("TEST4")
    
def test_guard_fails():
    """Shows how the function raises an exception when conditions are not met."""
    with pytest.raises(CustomGuardTestException):
        example_function("TEST")
    
    with pytest.raises(CustomGuardTestException):
        example_function(None)

These tests illustrate that the guard decorator prevents execution if the conditions are not satisfied, making precondition enforcement explicit.

Example Use Case

Below is a more practical example where the guard decorator ensures that a function receives a valid data connection:

@guard(  # Ensure a data connection is provided
    lambda _, data_connection: data_connection is not None,
    lambda: QueryEngineError("No data connection provided"),
)
@guard(  # Prevent the query engine from running on data connections that are not allowed
    lambda _, data_connection: data_connection["allowQueryEngine"] is True,
    lambda: QueryEngineError(
        "This data connection is not configured to be used by the Query Engine"
    ),
    raiseExceptionOnException=True,
)
def run_query(query: str, data_connection: dict) -> tuple:
    if query == "TEST":
        return [{"test": "test"}], ["test"]

    db_type_to_function = {
        "sqlite3": run_sqlite_query,
        "MySQL": run_mysql_query,
        "PostgreSQL": run_postgres_query,
        "MSSQL": run_mssql_query,
    }

    db_type = data_connection["resolvedParameters"]["type"]
    engine_query_function = db_type_to_function.get(db_type)

    if not engine_query_function:
        raise ValueError(f"Unsupported database type: {db_type}")

    return engine_query_function(query, data_connection)

Conclusion

The guard decorator brings a Swift-like precondition mechanism to Python. While not identical to Swift’s guard statement, it improves readability and maintainability by making preconditions explicit. This approach can be particularly useful when validating function inputs and ensuring constraints are met before execution.

Fitti Weissglas