wrapt + adapter which adds a argument combined with pytest yielding fixture

[jankatins]

I needed to wrap a pytest fixture (not test!) to let it only run once when running in xdist. For that I needed a adapter, which added a required argument (so that pytest would fetch the required tmp_path_factory fixture value and pass it in) and make it a generator (so pytest would see it as a yielding fixture, not a returning one).

The particular problem is that if I just manipulate the argspec, I end up with a fixture, which pytest identifies as "returning", not "yielding": pytest uses inspect.isgeneratorfunction() (which looks at some code specifics, not return types) to see if a function is a generator, but wrapt creates a function with just pass (=return None) and so gets identified as a non-Generator:

wrapt/src/wrapt/decorators.py

Line 196 in 24e2317

exec(f"def adapter{adapter}: pass", ns, ns)

It would be nice if wrapt would be able to generate a "proper" generator depending on some arguments (I guess async def will be equally problematic).

The code I came up with is below, it basically re-implements

wrapt/src/wrapt/decorators.py

Lines 177 to 204 in 24e2317

	if not callable(adapter):
	ns = {}
	# Check if the signature argument specification has
	# annotations. If it does then we need to remember
	# it but also drop it when attempting to manufacture
	# a standin adapter function. This is necessary else
	# it will try and look up any types referenced in
	# the annotations in the empty namespace we use,
	# which will fail.
	annotations = {}
	if not isinstance(adapter, str):
	if len(adapter) == 7:
	annotations = adapter[-1]
	adapter = adapter[:-1]
	adapter = formatargspec(*adapter)
	exec(f"def adapter{adapter}: pass", ns, ns)
	adapter = ns["adapter"]
	# Override the annotations for the manufactured
	# adapter function so they match the original
	# adapter signature argument specification.
	if annotations:
	adapter.__annotations__ = annotations

with the required generator signature:

RT = TypeVar("RT", bound=BaseModel)


def _additional_fixtures_protocol(tmp_path_factory: pytest.TempPathFactory):  # type: ignore[no-untyped-def]  # noqa: ANN202
    """Protocol to get access to the tmp_path_factory arg spec."""


def combine_args_with_protocol_adapter_factory(wrapped: Callable[..., Any]) -> Callable[..., Any]:
    """Adjust the signature of the wrapped functions with additional arguments from the protocol."""
    # At this point we know that the wrapped function is a fixture, so should only contain args.
    # We also know that the protocol only contains args

    argspec_wrapped = inspect.getfullargspec(wrapped)
    argspec_protocol = inspect.getfullargspec(_additional_fixtures_protocol)
    combined_args = argspec_wrapped.args[:] + argspec_protocol.args[:]

    adapter_spec = formatargspec(
        args=combined_args,
        varkw=argspec_wrapped.varkw,
        defaults=argspec_wrapped.defaults,
        kwonlyargs=argspec_wrapped.kwonlyargs,
        kwonlydefaults=argspec_wrapped.kwonlydefaults,
        varargs=argspec_wrapped.varargs,
        # No annotations, it would fail to compile
    )
    # the current wrapt produces a normal function, no yielding generator,
    # so we have to create one here ourselves with exec :-(
    # We need it because pytest uses inspect.isgeneratorfunction() to decide between generator
    # fixtures and fixtures with return and if we would have a return function, it would never get the
    # data out of the generator :-(
    ns: dict[str, Any] = {}
    exec_(f"def adapter{adapter_spec}: yield", ns, ns)
    adapter = ns["adapter"]
    # the protocol only contains arguments, no return type, so we would not override that...
    annotations = argspec_protocol.annotations.copy()
    # We prefer the annotations from the wrapped function, including the annotation for the return type
    annotations.update(argspec_wrapped.annotations.copy())
    adapter.__annotations__ = annotations
    return adapter


# Decorator modeled after
# https://pytest-xdist.readthedocs.io/en/latest/how-to.html#making-session-scoped-fixtures-execute-only-once
def xdist_run_only_once(  # noqa: PLR0915
    *, return_type: type[RT]
) -> Callable[[Callable[..., Iterator[RT]]], Callable[[pytest.TempPathFactory, str], Iterator[RT]]]:
    """Call a fixture only once despite xdist."""
    worker_id = os.environ.get("PYTEST_XDIST_WORKER", "master")
    if worker_id == "master":
        # not executing with multiple workers or without xdist
        # -> just make the decorator return the original functions
        return lambda x: x

    @wrapt.decorator(adapter=wrapt.adapter_factory(combine_args_with_protocol_adapter_factory))
    def adapted(  # type: ignore[no-untyped-def]  # noqa: PLR0915
        wrapped: Callable[..., Iterator[RT]],
        instance,  # noqa: ANN001
        # One of these already contains the new arguments from the protocol
        args,  # noqa: ANN001
        kwargs,  # noqa: ANN001
    ) -> Iterator[RT]:
        """Inner fixture with the interface of the combined arguments of the original fixture + the protocol."""
        lock_name = f"{wrapped.__module__}.{wrapped.__name__}"

        # The _executer function is a shorter way to pull out the named arguments from args/kwargs
        # no matter if these are in args or kwargs
        def _executer(  # type: ignore[no-untyped-def] # noqa: PLR0915
            tmp_path_factory: pytest.TempPathFactory,
            *_args,  # noqa: ANN002
            **_kwargs,  # noqa: ANN003
        ) -> Iterator[RT]:
            running_fixture: Iterator[Any] | None = None
            # get the temp directory shared by all workers
            # getbasetemp() is a worker specific directory under xdist, so go one down to get the shared one
            root_tmp_dir = tmp_path_factory.getbasetemp().parent

            lock_file = root_tmp_dir / f"{lock_name}.lock"
            info_file = root_tmp_dir / f"{lock_name}.json"
            worker_file = root_tmp_dir / f"{lock_name}.workers"

            def _load_worker_list() -> list[str]:
                if not worker_file.is_file():
                    return []
                return sorted(json.loads(worker_file.read_text())["workers"])

            def _write_worker_list(workers: list[str]) -> None:
                worker_file.write_text(json.dumps({"workers": workers}))

            def _add_worker() -> None:
                worker_id = os.environ["PYTEST_XDIST_WORKER"]
                with FileLock(str(lock_file)):
                    workers = _load_worker_list()
                    workers.append(worker_id)
                    _write_worker_list(workers)

            def _remove_worker() -> None:
                worker_id = os.environ["PYTEST_XDIST_WORKER"]
                with FileLock(str(lock_file)):
                    workers = _load_worker_list()
                    try:
                        workers.remove(worker_id)
                    except ValueError:

                    _write_worker_list(workers)

            with FileLock(str(lock_file)):
                if info_file.is_file():
                    data = return_type.model_validate_json(info_file.read_text())
                else:
                    # The first one actually creates it
                    running_fixture = wrapped(*_args, **_kwargs)
                    data = next(running_fixture)
                    info_file.write_text(data.model_dump_json())

            _add_worker()
            # Only yield when out of the locks!
            yield data
            _remove_worker()

            # We have nothing to do anymore, shut down any resources, but only if
            # - we created them and
            # - only after we are the last worker

            if running_fixture is None:
                return

            start = time.monotonic()
            timeout = 20 * 60  # 20 min
            # Wait for workers to become empty as other worker shut down
            while start + timeout > time.monotonic():
                with FileLock(str(lock_file)):
                    workers = _load_worker_list()
                if len(workers) == 0:
                    break
                time.sleep(1)
            # And now we are the last and can run the fixture clean up and then our own cleanup
            with FileLock(str(lock_file)):
                # We expect that ends with a raised StopIteration
                # BUT we have to return normally as otherwise this gets turned into a RuntimeError
                with contextlib.suppress(StopIteration):
                    next(running_fixture)
                worker_file.unlink(missing_ok=True)
                info_file.unlink(missing_ok=True)

        yield from _executer(*args, **kwargs)

    return adapted

[GrahamDumpleton]

Can you provide me the original code you were using when using the adapter function of @wrapt.decorator before introducing your special decorator factory? I want to see what you were originally passing as the adapter argument and how you were generating it.

[jankatins]

Good question, this is my last version in Intelij local history which didn't use the exec thingy:

def _xdist_run_only_once_protocol(tmp_path_factory: TempPathFactory, worker_id: str) -> Iterator[RT]:
    """Protocol of function which runs fixtures only once under xdist."""
    # We only need a yield so we get wrapt to expose the right magix flag to inspect.isgeneratorfunction(f)
    # to make pytest actually see this as a generator...
    yield ""  # type: ignore[misc]


def argspec_factory(wrapped) -> inspect.Arg:
    """Adjust the signature of the wrapped functions with additional arguments from the protocol"""
    argspec_wrapped = inspect.getfullargspec(wrapped)
    argspec_protocol = inspect.getfullargspec(_xdist_run_only_once_protocol)

    defaults = argspec_wrapped.defaults and argspec_wrapped.defaults[-len(argspec_wrapped.args):]

    # Now combine them
    args_wrapped = argspec_wrapped.args[:]
    args_protocol = argspec_protocol.args[:]
    combined_args = args_wrapped + args_protocol
    assert not argspec_wrapped.defaults
    assert not argspec_protocol.defaults

    return inspect.ArgSpec(combined_args, argspec_wrapped.varargs, argspec_wrapped.keywords, defaults, )

I think my first idea was going through the examples at https://wrapt.readthedocs.io/en/master/decorators.html#signature-changing-decorators and trying to return a version of return inspect.ArgSpec(args, argspec.varargs, argspec.keywords, defaults) which contained the additional arg. But that failed as pytest then didn't see a yielding fixture but a returning one.

[GrahamDumpleton]

The wrapt.adapter_factory is implemented as:

class AdapterFactory:
    def __call__(self, wrapped):
        raise NotImplementedError()


class DelegatedAdapterFactory(AdapterFactory):
    def __init__(self, factory):
        super(DelegatedAdapterFactory, self).__init__()
        self.factory = factory

    def __call__(self, wrapped):
        return self.factory(wrapped)


adapter_factory = DelegatedAdapterFactory

Best I can think of right now to try and fit in with how existing API is to have something like:

class AdapterFactory:
    def __init__(self, *, **options):
        self.options = options

    def __call__(self, wrapped):
        raise NotImplementedError()


class DelegatedAdapterFactory(AdapterFactory):
    def __init__(self, factory, *, **options):
        super(DelegatedAdapterFactory, self).__init__(**options)
        self.factory = factory

    def __call__(self, wrapped):
        return self.factory(wrapped)


adapter_factory = DelegatedAdapterFactory

and usage would be something like:

@wrapt.decorator(adapter=wrapt.adapter_factory(argspec_factory, body="yield"))

Implementation in @wrapt.decorator would be something like:

                body = "pass"
                if isinstance(adapter, AdapterFactory):
                    body = adapter.options.get("body") or "pass"
                    adapter = adapter(wrapped)

                if not callable(adapter):
                    ...

                    exec(f"def adapter{adapter}: {body}", ns, ns)

In other words, allow adapter factory to accept optional keyword parameters which can modify how adapter mechanism works. Allow substituting the body of the generated function would be one example. Question is whether should allow such low level injection, or instead use a boolean flag such as "generator" and have the code internally work out how that would be handled.

Have any thoughts on this suggestion? You still have to do all the arg spec manipulation yourself.

The other option is to provide alternative variants of adapter_factory function called something else, which map to different AdapterFactory implementations.

Technically, all the code:

                if not callable(adapter):
                    ns = {}

                    # Check if the signature argument specification has
                    # annotations. If it does then we need to remember
                    # it but also drop it when attempting to manufacture
                    # a standin adapter function. This is necessary else
                    # it will try and look up any types referenced in
                    # the annotations in the empty namespace we use,
                    # which will fail.

                    annotations = {}

                    if not isinstance(adapter, str):
                        if len(adapter) == 7:
                            annotations = adapter[-1]
                            adapter = adapter[:-1]
                        adapter = formatargspec(*adapter)

                    exec(f"def adapter{adapter}: pass", ns, ns)
                    adapter = ns["adapter"]

                    # Override the annotations for the manufactured
                    # adapter function so they match the original
                    # adapter signature argument specification.

                    if annotations:
                        adapter.__annotations__ = annotations

which currently sits in @wrapt.decorator could be within an instance of AdapterFactory code such that rather than the factory return either a callable or arg spec, it always returns a callable. That is, the factory would itself map the arg spec provided by your function to a callable rather than it being done in @wrapt.decorator code.

I'll have to think about best way to refactor it.

[jankatins]

Have any thoughts on this suggestion? You still have to do all the arg spec manipulation yourself.

Yeah, I don't think i can get around that.

In any case, these were also the things I came up with while reading the code:

• Make the adapter factory take a argument so this can be switched. (not sure if async def would also need such a argument)
• Make the adapter factory do this automatically by looking at the wrapped function

The latter make sense to me, as I would assume adding a decorator should not change this... but then: other ppl. do other interesting things and they might then need an explicit override, so I guess the first is still the more general approach...

Re moving the code into the adapter factory: reading the code that also came to my mind: it would give the adapter a clean interface: either pass in a function which represents the new signature/... OR pass in a factory which returns such a function when passed in the wrapped function. And include one which takes a signature returning function and turns that signature into a function as a nice convenience factory.

Another though was also that the adapter function would also have a need to decide if the resulting signature/adapter should be a generator or not (e.g. by querying the wrapped function). So maybe just returning the ArgSpec is not enough and it should also have a way to indicate that it should be a generator/async? No idea where to draw the line...

Anyway: and that's when it became all too complicated to me, and I stopped trying to come up with a PR idea, just implemented the above workaround, and wrote this ticket :-)