Refactor pytest fixtures

• Playground Link

Description

One of the most commonly used testing framework in Python is pytest. Among other things, it allows the use of fixtures.

Fixtures are defined as functions that can be required in test code, or in other fixtures, as an argument. This means that all functions arguments with a given name in a pytest context (test function or fixture) are essentially the same entity. However, not every editor's LSP is able to keep track of this, making refactoring challenging.

Using ast-grep, we can define some rules to match fixture definition and usage without catching similarly named entities in a non-test context.

First, we define utils to select pytest test/fixture functions.

utils:
  is-fixture-function:
    kind: function_definition
    follows:
      kind: decorator
      has:
        kind: identifier
        regex: ^fixture$
        stopBy: end
  is-test-function:
    kind: function_definition
    has:
      field: name
      regex: ^test_

Pytest fixtures are declared with a decorator @pytest.fixture. We match the function_definition node that directly follows a decorator node. That decorator node must have a fixture identifier somewhere. This accounts for different location of the fixture node depending on the type of imports and whether the decorator is used as is or called with parameters.

Pytest functions are fairly straghtforward to detect, as they always start with test_ by convention.

The next utils builds onto those two to incrementally:

• Find if a node is inside a pytest context (test/fixture)
• Find if a node is an argument in such a context

utils:
  is-pytest-context:
    # Pytest context is a node inside a pytest
    # test/fixture
    inside:
      stopBy: end
      any:
        - matches: is-fixture-function
        - matches: is-test-function
  is-fixture-arg:
    # Fixture arguments are identifiers inside the 
    # parameters of a test/fixture function
    all:
      - kind: identifier
      - inside:
          kind: parameters
      - matches: is-pytest-context

Once those utils are declared, you can perform various refactoring on a specific fixture.

The following rule adds a type-hint to a fixture.

rule:
  matches: is-fixture-arg
  regex: ^foo$
fix: 'foo: int'

This one renames a fixture and all its references.

rule:
  kind: identifier
  matches: is-fixture-context
  regex: ^foo$
fix: 'five'

Example

Renaming Fixtures

@pytest.fixture
def foo() -> int:
    return 5

@pytest.fixture(scope="function")
def some_fixture(foo: int) -> str:
    return str(foo)

def regular_function(foo) -> None:
    ...

def test_code(foo: int) -> None:
    assert foo == 5

Diff

@pytest.fixture
def foo() -> int: 
def five() -> int: 
    return 5

@pytest.fixture(scope="function")
def some_fixture(foo: int) -> str: 
def some_fixture(five: int) -> str: 
    return str(foo)

def regular_function(foo) -> None:
    ...

def test_code(foo: int) -> None: 
def test_code(five: int) -> None: 
    assert foo == 5
    assert five == 5

Type Hinting Fixtures

@pytest.fixture
def foo() -> int:
    return 5

@pytest.fixture(scope="function")
def some_fixture(foo) -> str:
    return str(foo)

def regular_function(foo) -> None:
    ...

def test_code(foo) -> None:
    assert foo == 5

Diff

@pytest.fixture
def foo() -> int:
    return 5

@pytest.fixture(scope="function")
def some_fixture(foo) -> str: 
def some_fixture(foo: int) -> str: 
    return str(foo)

def regular_function(foo) -> None:
    ...

def test_code(foo) -> None: 
def test_code(foo: int) -> None: 
    assert foo == 5