codekingpro/portable-devtools

1# Copyright The OpenTelemetry Authors2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15import contextlib16import functools17import inspect18from typing import TYPE_CHECKING, Callable, Generic, Iterator, TypeVar19 20V = TypeVar("V")21R = TypeVar("R")  # Return type22Pargs = TypeVar("Pargs")  # Generic type for arguments23Pkwargs = TypeVar("Pkwargs")  # Generic type for arguments24 25# We don't actually depend on typing_extensions but we can use it in CI with this conditional26# import. ParamSpec can be imported directly from typing after python 3.9 is dropped27# https://peps.python.org/pep-0612/.28if TYPE_CHECKING:29    from typing_extensions import ParamSpec30 31    P = ParamSpec("P")  # Generic type for all arguments32 33 34class _AgnosticContextManager(35    contextlib._GeneratorContextManager[R],36    Generic[R],37):  # pylint: disable=protected-access38    """Context manager that can decorate both async and sync functions.39 40    This is an overridden version of the contextlib._GeneratorContextManager41    class that will decorate async functions with an async context manager42    to end the span AFTER the entire async function coroutine finishes.43 44    Else it will report near zero spans durations for async functions.45 46    We are overriding the contextlib._GeneratorContextManager class as47    reimplementing it is a lot of code to maintain and this class (even if it's48    marked as protected) doesn't seems like to be evolving a lot.49 50    For more information, see:51    https://github.com/open-telemetry/opentelemetry-python/pull/363352    """53 54    def __enter__(self) -> R:55        """Reimplementing __enter__ to avoid the type error.56 57        The original __enter__ method returns Any type, but we want to return R.58        """59        del self.args, self.kwds, self.func  # type: ignore60        try:61            return next(self.gen)  # type: ignore62        except StopIteration:63            raise RuntimeError("generator didn't yield") from None64 65    def __call__(self, func: V) -> V:  # pyright: ignore [reportIncompatibleMethodOverride]66        if inspect.iscoroutinefunction(func):67 68            @functools.wraps(func)  # type: ignore69            async def async_wrapper(*args: Pargs, **kwargs: Pkwargs) -> R:  # pyright: ignore [reportInvalidTypeVarUse]70                with self._recreate_cm():  # type: ignore71                    return await func(*args, **kwargs)  # type: ignore72 73            return async_wrapper  # type: ignore74        return super().__call__(func)  # type: ignore75 76 77def _agnosticcontextmanager(78    func: "Callable[P, Iterator[R]]",79) -> "Callable[P, _AgnosticContextManager[R]]":80    @functools.wraps(func)81    def helper(*args: Pargs, **kwargs: Pkwargs) -> _AgnosticContextManager[R]:  # pyright: ignore [reportInvalidTypeVarUse]82        return _AgnosticContextManager(func, args, kwargs)  # pyright: ignore [reportArgumentType]83 84    # Ignoring the type to keep the original signature of the function85    return helper  # type: ignore[return-value]86