Supercharging Python Automation with Nox: Beyond Basic Usage

Learn how to extend Nox with custom decorators and classes to create a more flexible, powerful automation system for your Python projects.

A deep dive into extending Nox’s capabilities with custom decorators and classes, turning it into a comprehensive automation solution that can replace Makefiles and shell scripts.

Introduction

Python automation tools are essential for maintaining consistency and efficiency in development workflows. Recently, I came across the Nox library through a YouTube video, and it transformed how I handle automation in my Python projects. Unlike traditional tools like Makefiles or shell scripts, Nox offers a more Pythonic, flexible, and readable approach to automation tasks. This post explores how I extended Nox’s functionality with custom decorators and classes to make it even more powerful.

Understanding Nox

Nox is a command-line tool that automates testing in multiple Python environments, similar to tox but with a more straightforward, Python-based configuration. Instead of using a configuration file like tox.ini, Nox uses a Python file (noxfile.py) where you define sessions as Python functions. This approach gives you the full power of Python for configuring your automation tasks.

The core concept in Nox is a “session,” which is essentially a function decorated with @nox.session that defines a discrete task such as running tests, building documentation, or deploying code. Each session typically creates a virtual environment, installs dependencies, and then runs commands within that environment.

Extending Nox with Custom Functionality

While Nox provides excellent baseline functionality, I found myself repeatedly implementing certain patterns across projects. To streamline this process, I developed extensions to Nox that address common needs:

Installing dependency groups from pyproject.toml
Managing environment variables dynamically
Providing default session arguments
Temporarily altering session behavior

Let’s explore each of these extensions in detail.

The Custom Session Decorator

The standard Nox session decorator is powerful, but I wanted to extend it with additional capabilities. My custom @session decorator builds on Nox’s foundation while adding support for dependency groups, environment variables, and default arguments.

def session(
    f: "Callable[..., Any]" = None,
    /,
    dependency_group: str = None,
    environment_mapping: "Dict[str, str]" = {},
    default_posargs: "Sequence[str]" = (),
    **kwargs: "NoxSessionParams",
) -> "Callable[..., Any]":
    if f is None:
        # so that `@session` can be used like so:
        # ```python
        # @session
        # def some_task(...):
        #     ...
        # ```
        return lambda f: session(
            f,
            dependency_group=dependency_group,
            environment_mapping=environment_mapping,
            default_posargs=default_posargs,
            **kwargs,
        )
    f_name = f.__name__.replace("_", "-")
    nox_session_kwargs = {
        **DEFAULT_SESSION_KWARGS,
        **kwargs,
    }
    nox_session_kwargs["name"] = f_name

    @wraps(f)
    def wrapper(session: Session, *args, **kwargs):
        altered_session = AlteredSession(
            session, dependency_group, environment_mapping, default_posargs
        )
        return f(altered_session, *args, **kwargs)

    return nox_session(wrapper, **nox_session_kwargs)

This decorator works by accepting either a function directly or acting as a factory when called without a function argument. When creating a session, it converts underscores in the function name to hyphens (e.g. fastapi_auth as function name to fastapi-auth as Nox’s CLI’s session name), merges default session parameters with any provided overrides, and wraps the original function to use an enhanced session object.

The key innovation is the creation of an AlteredSession instance that extends the standard Nox Session with additional capabilities. This approach preserves all standard Nox functionality while seamlessly adding new features.

The AlteredSession Class

The AlteredSession class extends Nox’s Session class to add support for dependency groups, environment variables, and default arguments:

class AlteredSession(Session):
    __slots__ = (
        "session",
        "dependency_group",
        "environment_mapping",
        "default_posargs",
    )

    def __init__(
        self,
        session: Session,
        dependency_group: str,
        environment_mapping: "Dict[str, str]",
        default_posargs: "Sequence[str]",
    ):
        super().__init__(session._runner)
        self.dependency_group = dependency_group
        self.environment_mapping = environment_mapping
        self.default_posargs = default_posargs
        self.session = session

    def run(self, *args, **kwargs):
        if self.dependency_group is not None:
            uv_install_group_dependencies(self, self.dependency_group)
        if self.session.posargs is not None:
            args = (*args, *(self.session.posargs or self.default_posargs))
        env: "Dict[str, str]" = kwargs.pop("env", {})
        env.update(self.environment_mapping)
        kwargs["env"] = env
        return self.session.run(*args, **kwargs)

This class overrides the run method to provide several enhancements. First, if a dependency group is specified, it automatically installs those dependencies using the uv_install_group_dependencies function. Second, it appends any positional arguments passed to the Nox command line (or default arguments if none were provided). Finally, it updates the environment variables for the command being run with any specified in the environment mapping.

The dependency group functionality relies on a helper function that parses the pyproject.toml file:

def uv_install_group_dependencies(session: Session, dependency_group: str):
    pyproject = nox.project.load_toml(MANIFEST_FILENAME)
    dependencies = nox.project.dependency_groups(pyproject, dependency_group)
    session.install(*dependencies)
    session.log(f"Installed dependencies: {dependencies} for {dependency_group}")

This function uses Nox’s built-in project utilities to load the pyproject.toml file and extract the dependencies for a specific group. It then installs these dependencies and logs the process for transparency.

The alter_session Context Manager

While the enhanced session decorator is useful for setting up default behavior, sometimes you need to temporarily modify session parameters for specific tasks. The alter_session context manager provides this flexibility:

@contextlib.contextmanager
def alter_session(
    session: AlteredSession,
    dependency_group: str = None,
    environment_mapping: "Dict[str, str]" = {},
    default_posargs: "Sequence[str]" = (),
    **kwargs: "NoxSessionParams",
):
    old_dependency_group = session.dependency_group
    old_environment_mapping = session.environment_mapping
    old_default_posargs = session.default_posargs
    old_kwargs = {}
    for key, value in kwargs.items():
        old_kwargs[key] = getattr(session, key)

    session.dependency_group = dependency_group
    session.environment_mapping = environment_mapping
    session.default_posargs = default_posargs
    for key, value in kwargs.items():
        setattr(session, key, value)
    yield session

    session.dependency_group = old_dependency_group
    session.environment_mapping = old_environment_mapping
    session.default_posargs = old_default_posargs
    for key, value in old_kwargs.items():
        setattr(session, key, value)

This context manager temporarily modifies an AlteredSession instance’s properties, yields control back to the caller, and then restores the original properties when the context block exits. This pattern is particularly useful for running the same function with different configurations without duplicating code.

Practical Examples

Let’s explore some practical examples of using these extensions in real-world scenarios.

Running Tests with Default Arguments

One common automation task is running tests. With the enhanced session decorator, you can define default arguments while still allowing command-line overrides:

@session(dependency_group="test", default_posargs=["tests", "-s", "-vv"])
def test(session: AlteredSession):
    command = [
        shutil.which("uv"),
        "run",
        "--group",
        "test",
        "python",
        "-m",
        "pytest",
    ]
    if "--build" in session.posargs:
        session.posargs.remove("--build")
        with alter_session(session, dependency_group="build"):
            build(session)
    session.run(*command)

This session function does several clever things. First, it specifies a “test” dependency group, automatically installing all test dependencies. Second, it sets default positional arguments that will be used if none are provided on the command line. Finally, it checks for a --build flag in the arguments and runs a build session first if it’s present, using the alter_session context manager to temporarily switch to the “build” dependency group.

Running this session is as simple as nox -s test, which will run all tests with verbose output. Alternatively, you can specify a particular test file with nox -s test -- tests/test_specific.py, and the session will use your arguments instead of the defaults.

Testing Multiple Environments

Another powerful use case is testing an application across multiple environments, such as development, staging, and production:

@session(
    dependency_group="examples",
)
def fastapi_auth(session: Session):
    test_development(session)
    # test the staging environment
    # change the environment key to "staging"
    with alter_session(session, environment_mapping={"ENVIRONMENT_KEY": "staging"}):
        test_staging(session)
    # test the production environment
    # change the environment key to "production"
    with alter_session(session, environment_mapping={"ENVIRONMENT_KEY": "production"}):
        test_production(session)
    # test the development environment again with the environment key set to "development"
    test_development(session)

This example demonstrates how the alter_session context manager makes it easy to run the same tests against different environments by simply changing environment variables. Each environment’s tests run in isolation with the appropriate configuration, without needing to create separate session functions for each environment.

Advantages Over Traditional Approaches

This extended Nox approach offers several advantages over traditional automation methods:

First, it provides a more intuitive, Pythonic syntax compared to Makefiles or shell scripts. The code is easier to read, write, and maintain, especially for developers already familiar with Python.

Second, the dependency group support ensures that the right dependencies are installed for each task, eliminating the need for separate environment setup steps or complex conditional logic in your automation scripts.

Third, the ability to temporally modify session parameters with the context manager reduces code duplication and makes it easier to run variations of the same task with different configurations.

Fourth, the extension maintains compatibility with Nox’s core functionality, allowing you to leverage Nox’s existing features while adding new capabilities.

Fifth, the flexibility of Python means you can incorporate complex logic, conditional execution, and dynamic configuration in your automation tasks, going beyond what’s easily achievable with Makefiles or shell scripts.

Conclusion

Nox offers a powerful foundation for Python automation, and with a few custom extensions, it can replace Makefiles and shell scripts entirely in many Python projects. By creating a custom session decorator, an extended session class, and a context manager for temporarily modifying session parameters, you can build a more flexible, maintainable automation system that adapts to your specific needs.

The approach demonstrated in this post not only makes automation more Pythonic but also addresses common pain points like dependency management, environment configuration, and command-line argument handling. Whether you’re managing a small personal project or a large, complex application, these Nox extensions can simplify your workflow and improve your development experience.

If you’re still using Makefiles or shell scripts for Python project automation, I encourage you to give Nox a try, perhaps with the extensions described in this post. Your future self will thank you for the improved readability, flexibility, and maintainability of your automation code.