Separate setup, exercise, and validate in pytest with hypothesis as coverage contexts

Question

In our code base, we have a bunch of property-based unit tests. We would like to get the context-annotated code coverage of our unit tests. Pytest can do this, but it needs to know what part is the Run (aka Exercise, Act) phase as opposed to Setup/Arrange. How do we mark the boundary if we use Hypothesis?

For example, we have this unit test.

@dataclasses.dataclass
class StepDummy:
    """A dummy object pretending to be a Step object for typing purposes."""

    step_id: int = -1

@given(
    n_steps=st.lists(st.integers(0, 10), min_size=0, max_size=10),
    n_steps_to_remove=st.integers(0, 100),
)
def test_remove_first_n_steps_from_tasks(n_steps: list[int], n_steps_to_remove: int) -> None:
    """Test that remove_first_n_steps removes exactly n steps."""
    # ARRANGE
    tasks = [
        Task(
            task_id=m,
            steps=[
                cast(Step, StepDummy(i))  # We don't care about the internals of a step
                for i in range(n_steps_in_task)
            ],
        )
        for m, n_steps_in_task in enumerate(n_steps)
    ]
    total_n_steps_before = sum(len(m.steps) for m in tasks)
    assume(n_steps_to_remove <= total_n_steps_before)
    
    # ACT
    reduced_tasks = remove_first_n_steps_from_tasks(tasks, n_steps_to_remove)

    # ASSERT
    total_n_steps_after = sum(len(m.steps) for m in reduced_tasks)
    assert total_n_steps_after + n_steps_to_remove == total_n_steps_before

Without Hypothesis, I could write it like this.

@pytest
def test_remove_first_n_steps_from_tasks_setup() -> None:
    """Setup for ."""
    n_steps = [7, 10]
    n_steps_to_remove = 8
    tasks = [
        Task(
            task_id=m,
            steps=[
                cast(Step, StepDummy(i))  # We don't care about the internals of a step
                for i in range(n_steps_in_task)
            ],
        )
        for m, n_steps_in_task in enumerate(n_steps)
    ]
    total_n_steps_before = sum(len(m.steps) for m in tasks)
    return tasks, n_steps_to_remove, total_n_steps_before

def test_remove_first_n_steps_from_tasks(test_remove_first_n_steps_from_tasks_setup) -> None:
    """Test that remove_first_n_steps removes exactly n steps."""
    (
        tasks, n_steps_to_remove, total_n_steps_before
    ) = test_remove_first_n_steps_from_tasks_setup
    # ACT
    reduced_tasks = remove_first_n_steps_from_tasks(tasks, n_steps_to_remove)

    # ASSERT
    total_n_steps_after = sum(len(m.steps) for m in reduced_tasks)
    assert total_n_steps_after + n_steps_to_remove == total_n_steps_before

Not particularly pretty, but possible. How do I do this with hypothesis? Is there maybe even a way to set the pytest coverage dynamic context to setup or run manually?