# Source code for python_utils.decorators

    import contextlib
    import functools
    import logging
    import random
    from . import types

    _T = types.TypeVar('_T')
    _TC = types.TypeVar('_TC', bound=types.Container[types.Any])
    _P = types.ParamSpec('_P')
    _S = types.TypeVar('_S', covariant=True)

    def set_attributes(**kwargs: types.Any) -> types.Callable[..., types.Any]:
        '''Decorator to set attributes on functions and classes

        A common usage for this pattern is the Django Admin where
        functions can get an optional short_description. To illustrate:

        Example from the Django admin using this decorator:
        https://docs.djangoproject.com/en/3.0/ref/contrib/admin/#django.contrib.admin.ModelAdmin.list_display

        Our simplified version:

        >>> @set_attributes(short_description='Name')
        ... def upper_case_name(self, obj):
        ...     return ("%s %s" % (obj.first_name, obj.last_name)).upper()

        The standard Django version:

        >>> def upper_case_name(obj):
        ...     return ("%s %s" % (obj.first_name, obj.last_name)).upper()

        >>> upper_case_name.short_description = 'Name'

        '''

        def _set_attributes(
            function: types.Callable[_P, _T]
        ) -> types.Callable[_P, _T]:
            for key, value in kwargs.items():
                setattr(function, key, value)
            return function

        return _set_attributes

    def listify(
        collection: types.Callable[
            [types.Iterable[_T]], _TC
        ] = list,  # type: ignore
        allow_empty: bool = True,
    ) -> types.Callable[
        [types.Callable[..., types.Optional[types.Iterable[_T]]]],
        types.Callable[..., _TC],
    ]:
        '''
        Convert any generator to a list or other type of collection.

        >>> @listify()
        ... def generator():
        ...     yield 1
        ...     yield 2
        ...     yield 3

        >>> generator()
        [1, 2, 3]

        >>> @listify()
        ... def empty_generator():
        ...     pass

        >>> empty_generator()
        []

        >>> @listify(allow_empty=False)
        ... def empty_generator_not_allowed():
        ...     pass

        >>> empty_generator_not_allowed()  # doctest: +ELLIPSIS
        Traceback (most recent call last):
        ...
        TypeError: ... `allow_empty` is `False`

        >>> @listify(collection=set)
        ... def set_generator():
        ...     yield 1
        ...     yield 1
        ...     yield 2

        >>> set_generator()
        {1, 2}

        >>> @listify(collection=dict)
        ... def dict_generator():
        ...     yield 'a', 1
        ...     yield 'b', 2

        >>> dict_generator()
        {'a': 1, 'b': 2}
        '''

        def _listify(
            function: types.Callable[..., types.Optional[types.Iterable[_T]]]
        ) -> types.Callable[..., _TC]:
            def __listify(*args: types.Any, **kwargs: types.Any) -> _TC:
                result: types.Optional[types.Iterable[_T]] = function(
                    *args, **kwargs
                )
                if result is None:
                    if allow_empty:
                        return collection(iter(()))
                    else:
                        raise TypeError(
                            f'{function} returned `None` and `allow_empty` '
                            'is `False`'
                        )
                else:
                    return collection(result)

            return __listify

        return _listify

    def sample(sample_rate: float):
        '''
        Limit calls to a function based on given sample rate.
        Number of calls to the function will be roughly equal to
        sample_rate percentage.

        Usage:

        >>> @sample(0.5)
        ... def demo_function(*args, **kwargs):
        ...     return 1

        Calls to *demo_function* will be limited to 50% approximatly.
        '''

        def _sample(
            function: types.Callable[_P, _T]
        ) -> types.Callable[_P, types.Optional[_T]]:
            @functools.wraps(function)
            def __sample(
                *args: _P.args, **kwargs: _P.kwargs
            ) -> types.Optional[_T]:
                if random.random() < sample_rate:
                    return function(*args, **kwargs)
                else:
                    logging.debug(
                        'Skipped execution of %r(%r, %r) due to sampling',
                        function,
                        args,
                        kwargs,
                    )  # noqa: E501
                    return None

            return __sample

        return _sample

    def wraps_classmethod(
        wrapped: types.Callable[types.Concatenate[_S, _P], _T],
    ) -> types.Callable[
        [
            types.Callable[types.Concatenate[types.Any, _P], _T],
        ],
        types.Callable[types.Concatenate[types.Type[_S], _P], _T],
    ]:
        '''
        Like `functools.wraps`, but for wrapping classmethods with the type info
        from a regular method
        '''

        def _wraps_classmethod(
            wrapper: types.Callable[types.Concatenate[types.Any, _P], _T],
        ) -> types.Callable[types.Concatenate[types.Type[_S], _P], _T]:
            # For some reason `functools.update_wrapper` fails on some test
            # runs but not while running actual code
            with contextlib.suppress(AttributeError):
                wrapper = functools.update_wrapper(
                    wrapper,
                    wrapped,
                    assigned=tuple(
                        a
                        for a in functools.WRAPPER_ASSIGNMENTS
                        if a != '__annotations__'
                    ),
                )
            if annotations := getattr(wrapped, '__annotations__', {}):
                annotations.pop('self', None)
                wrapper.__annotations__ = annotations

            return wrapper

        return _wraps_classmethod