python-utils 2.5.6__py3-none-any.whl → 4.0.0__py3-none-any.whl

python_utils/decorators.py

@@ -1,7 +1,34 @@
+"""
+This module provides various utility decorators for Python functions
+and methods.
 
+The decorators include:
 
-def set_attributes(**kwargs):
-    '''Decorator to set attributes on functions and classes
+1. `set_attributes`: Sets attributes on functions and classes.
+2. `listify`: Converts any generator to a list or other collection.
+3. `sample`: Limits calls to a function based on a sample rate.
+4. `wraps_classmethod`: Wraps classmethods with type info from a
+   regular method.
+
+Each decorator is designed to enhance the functionality of Python
+functions and methods in a simple and reusable manner.
+"""
+
+import collections.abc
+import contextlib
+import functools
+import logging
+import random
+import typing
+
+_T = typing.TypeVar('_T')
+_P = typing.ParamSpec('_P')
+
+
+def set_attributes(
+    **kwargs: typing.Any,
+) -> collections.abc.Callable[..., typing.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:
@@ -13,18 +40,201 @@ def set_attributes(**kwargs):
 
     >>> @set_attributes(short_description='Name')
     ... def upper_case_name(self, obj):
-    ...     return ("%s %s" % (obj.first_name, obj.last_name)).upper()
+    ...     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()
+    ...     return ('%s %s' % (obj.first_name, obj.last_name)).upper()
 
     >>> upper_case_name.short_description = 'Name'
 
-    '''
-    def _set_attributes(function):
+    """
+
+    def _set_attributes(
+        function: collections.abc.Callable[_P, _T],
+    ) -> collections.abc.Callable[_P, _T]:
+        """Attach the captured ``kwargs`` as attributes on ``function``."""
         for key, value in kwargs.items():
             setattr(function, key, value)
         return function
+
     return _set_attributes
+
+
+def listify(
+    collection: collections.abc.Callable[
+        [collections.abc.Iterable[_T]], collections.abc.Collection[_T]
+    ] = list,
+    allow_empty: bool = True,
+) -> collections.abc.Callable[
+    [collections.abc.Callable[..., collections.abc.Iterable[_T] | None]],
+    collections.abc.Callable[..., collections.abc.Collection[_T]],
+]:
+    """
+    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: collections.abc.Callable[
+            ..., collections.abc.Iterable[_T] | None
+        ],
+    ) -> collections.abc.Callable[..., collections.abc.Collection[_T]]:
+        """Materialize ``function``'s result into ``collection``."""
+
+        def __listify(
+            *args: typing.Any, **kwargs: typing.Any
+        ) -> collections.abc.Collection[_T]:
+            """Call ``function`` and gather its result into ``collection``."""
+            result: collections.abc.Iterable[_T] | None = 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,
+) -> collections.abc.Callable[
+    [collections.abc.Callable[_P, _T]],
+    collections.abc.Callable[_P, _T | None],
+]:
+    """
+    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% approximately.
+    """
+
+    def _sample(
+        function: collections.abc.Callable[_P, _T],
+    ) -> collections.abc.Callable[_P, _T | None]:
+        """Wrap ``function`` so it only runs on a sampled fraction of calls."""
+
+        @functools.wraps(function)
+        def __sample(*args: _P.args, **kwargs: _P.kwargs) -> _T | None:
+            """Run ``function`` with probability ``sample_rate``, else skip."""
+            if random.random() < sample_rate:
+                return function(*args, **kwargs)
+            else:
+                logging.debug(
+                    'Skipped execution of %r(%r, %r) due to sampling',
+                    function,
+                    args,
+                    kwargs,
+                )
+                return None
+
+        return __sample
+
+    return _sample
+
+
+def wraps_classmethod(
+    wrapped: collections.abc.Callable[typing.Concatenate[typing.Any, _P], _T],
+) -> collections.abc.Callable[
+    [
+        collections.abc.Callable[typing.Concatenate[typing.Any, _P], _T],
+    ],
+    collections.abc.Callable[typing.Concatenate[typing.Any, _P], _T],
+]:
+    """Like ``functools.wraps``, but for wrapping classmethods.
+
+    Copies the wrapped method's metadata (name, docstring and annotations) onto
+    the wrapper, so a classmethod wrapper carries the type information of the
+    regular method it stands in for.
+
+    Args:
+        wrapped: The method whose metadata should be copied onto the wrapper.
+
+    Returns:
+        A decorator that updates its wrapper with ``wrapped``'s metadata.
+    """
+
+    def _wraps_classmethod(
+        wrapper: collections.abc.Callable[
+            typing.Concatenate[typing.Any, _P], _T
+        ],
+    ) -> collections.abc.Callable[typing.Concatenate[typing.Any, _P], _T]:
+        """Copy ``wrapped``'s metadata onto ``wrapper`` and return it."""
+        # 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__', {}):
+            # Drop `self`: the wrapper is a classmethod, so it takes no `self`.
+            annotations.pop('self', None)
+            wrapper.__annotations__ = annotations
+
+        return wrapper
+
+    return _wraps_classmethod

python_utils/exceptions.py

@@ -0,0 +1,47 @@
+"""
+This module provides utility functions for raising and reraising exceptions.
+
+Functions::
+
+    raise_exception(exception_class, *args, **kwargs):
+        Returns a function that raises an exception of the given type with
+        the given arguments.
+
+    reraise(*args, **kwargs):
+        Reraises the current exception.
+"""
+
+import collections.abc
+import typing
+
+
+def raise_exception(
+    exception_class: type[Exception],
+    *args: typing.Any,
+    **kwargs: typing.Any,
+) -> collections.abc.Callable[..., None]:
+    """
+    Returns a function that raises an exception of the given type with the
+    given arguments.
+
+    >>> raise_exception(ValueError, 'spam')('eggs')
+    Traceback (most recent call last):
+        ...
+    ValueError: spam
+    """
+
+    def raise_(*args_: typing.Any, **kwargs_: typing.Any) -> typing.Any:
+        """Raise ``exception_class`` with the captured args."""
+        raise exception_class(*args, **kwargs)
+
+    return raise_
+
+
+def reraise(*args: typing.Any, **kwargs: typing.Any) -> typing.Any:
+    """
+    Reraises the current exception.
+
+    This function seems useless, but it can be useful when you need to pass
+    a callable to another function that raises an exception.
+    """
+    raise

python_utils/formatters.py

@@ -1,8 +1,23 @@
+"""
+This module provides utility functions for formatting strings and dates.
+
+Functions::
+
+    camel_to_underscore: Convert camel case naming to underscore/snake case.
+    apply_recursive: Apply a function to all keys in a scope recursively.
+    timesince: Return a string representing 'time since', e.g. 3 days ago.
+"""
+
+# pyright: reportUnnecessaryIsInstance=false
+import collections.abc
 import datetime
+import typing
 
+from python_utils import _aliases
 
-def camel_to_underscore(name):
-    '''Convert camel case style naming to underscore style naming
+
+def camel_to_underscore(name: str) -> str:
+    """Convert camel case style naming to underscore/snake case style naming.
 
     If there are existing underscores they will be collapsed with the
     to-be-added underscores. Multiple consecutive capital letters will not be
@@ -18,8 +33,8 @@ def camel_to_underscore(name):
     '__spam_and_bacon__'
     >>> camel_to_underscore('__SpamANDBacon__')
     '__spam_and_bacon__'
-    '''
-    output = []
+    """
+    output: list[str] = []
     for i, c in enumerate(name):
         if i > 0:
             pc = name[i - 1]
@@ -30,7 +45,7 @@ def camel_to_underscore(name):
             elif i > 3 and not c.isupper():
                 # Will return the last 3 letters to check if we are changing
                 # case
-                previous = name[i - 3:i]
+                previous = name[i - 3 : i]
                 if previous.isalpha() and previous.isupper():
                     output.insert(len(output) - 1, '_')
 
@@ -39,8 +54,50 @@ def camel_to_underscore(name):
     return ''.join(output)
 
 
-def timesince(dt, default='just now'):
-    '''
+def apply_recursive(
+    function: collections.abc.Callable[[str], str],
+    data: _aliases.OptionalScope = None,
+    **kwargs: typing.Any,
+) -> _aliases.OptionalScope:
+    """
+    Apply a function to all keys in a scope recursively.
+
+    >>> apply_recursive(camel_to_underscore, {'SpamEggsAndBacon': 'spam'})
+    {'spam_eggs_and_bacon': 'spam'}
+    >>> apply_recursive(
+    ...     camel_to_underscore,
+    ...     {
+    ...         'SpamEggsAndBacon': {
+    ...             'SpamEggsAndBacon': 'spam',
+    ...         }
+    ...     },
+    ... )
+    {'spam_eggs_and_bacon': {'spam_eggs_and_bacon': 'spam'}}
+
+    >>> a = {'a_b_c': 123, 'def': {'DeF': 456}}
+    >>> b = apply_recursive(camel_to_underscore, a)
+    >>> b
+    {'a_b_c': 123, 'def': {'de_f': 456}}
+
+    >>> apply_recursive(camel_to_underscore, None)
+    """
+    if data is None:
+        return None
+
+    elif isinstance(data, dict):
+        return {
+            function(key): apply_recursive(function, value, **kwargs)
+            for key, value in data.items()
+        }
+    else:
+        return data
+
+
+def timesince(
+    dt: datetime.datetime | datetime.timedelta,
+    default: str = 'just now',
+) -> str:
+    """
     Returns string representing 'time since' e.g.
     3 days ago, 5 hours ago etc.
 
@@ -81,7 +138,7 @@ def timesince(dt, default='just now'):
     '1 hour and 2 minutes ago'
     >>> timesince(datetime.timedelta(seconds=3721))
     '1 hour and 2 minutes ago'
-    '''
+    """
     if isinstance(dt, datetime.timedelta):
         diff = dt
     else:
@@ -98,16 +155,15 @@ def timesince(dt, default='just now'):
         (diff.seconds % 60, 'second', 'seconds'),
     )
 
-    output = []
+    output: list[str] = []
     for period, singular, plural in periods:
-        if int(period):
-            if int(period) == 1:
-                output.append('%d %s' % (period, singular))
-            else:
-                output.append('%d %s' % (period, plural))
+        int_period = int(period)
+        if int_period == 1:
+            output.append(f'{int_period} {singular}')
+        elif int_period:
+            output.append(f'{int_period} {plural}')
 
     if output:
-        return '%s ago' % ' and '.join(output[:2])
+        return f'{" and ".join(output[:2])} ago'
 
     return default
-

python_utils/generators.py

@@ -0,0 +1,126 @@
+"""
+This module provides generator utilities for batching items from
+iterables and async iterables.
+
+Functions:
+    abatcher(generator, batch_size=None, interval=None):
+        Asyncio generator wrapper that returns items with a given batch
+        size or interval (whichever is reached first).
+
+    batcher(iterable, batch_size=10):
+        Generator wrapper that returns items with a given batch size.
+"""
+
+import asyncio
+import collections.abc
+import time
+import typing
+
+import python_utils
+from python_utils import _aliases
+
+#: Element type of the iterables being batched.
+_T = typing.TypeVar('_T')
+
+
+async def abatcher(
+    generator: collections.abc.AsyncGenerator[_T, None]
+    | collections.abc.AsyncIterator[_T],
+    batch_size: int | None = None,
+    interval: _aliases.delta_type | None = None,
+) -> collections.abc.AsyncGenerator[list[_T], None]:
+    """
+    Asyncio generator wrapper that returns items with a given batch size or
+    interval (whichever is reached first).
+
+    Args:
+        generator: The async generator or iterator to batch.
+        batch_size (typing.Optional[int], optional): The number of items per
+            batch. Defaults to None.
+        interval (typing.Optional[_aliases.delta_type], optional): The time
+            interval to wait before yielding a batch. Defaults to None.
+
+    Yields:
+        collections.abc.AsyncGenerator[list[_T], None]: A generator that yields
+        batches of items.
+    """
+    batch: list[_T] = []
+
+    assert batch_size or interval, 'Must specify either batch_size or interval'
+
+    # If interval is specified, use it to determine when to yield the batch
+    # Alternatively set a really long timeout to keep the code simpler
+    if interval:
+        interval_s = python_utils.delta_to_seconds(interval)
+    else:
+        # Set the timeout to 10 years
+        interval_s = 60 * 60 * 24 * 365 * 10.0
+
+    next_yield: float = time.perf_counter() + interval_s
+
+    done: set[asyncio.Task[_T]]
+    pending: set[asyncio.Task[_T]] = set()
+
+    while True:
+        try:
+            done, pending = await asyncio.wait(
+                pending
+                or [
+                    asyncio.create_task(
+                        typing.cast(
+                            collections.abc.Coroutine[None, None, _T],
+                            generator.__anext__(),
+                        )
+                    ),
+                ],
+                timeout=interval_s,
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+
+            if done:
+                batch.extend(result.result() for result in done)
+
+        except StopAsyncIteration:
+            if batch:
+                yield batch
+
+            break
+
+        if batch_size is not None and len(batch) == batch_size:
+            yield batch
+            batch = []
+
+        if interval and batch and time.perf_counter() > next_yield:
+            yield batch
+            batch = []
+            # Always set the next yield time to the current time. If the
+            # loop is running slow due to blocking functions we do not
+            # want to burst too much
+            next_yield = time.perf_counter() + interval_s
+
+
+def batcher(
+    iterable: collections.abc.Iterable[_T],
+    batch_size: int = 10,
+) -> collections.abc.Generator[list[_T], None, None]:
+    """
+    Generator wrapper that returns items with a given batch size.
+
+    Args:
+        iterable (collections.abc.Iterable[_T]): The iterable to batch.
+        batch_size (int, optional): The number of items per batch. Defaults
+            to 10.
+
+    Yields:
+        collections.abc.Generator[list[_T], None, None]: A generator that
+            yields batches of items.
+    """
+    batch: list[_T] = []
+    for item in iterable:
+        batch.append(item)
+        if len(batch) == batch_size:
+            yield batch
+            batch = []
+
+    if batch:
+        yield batch

python_utils/import_.py

@@ -1,32 +1,62 @@
+"""
+This module provides utilities for importing modules and handling exceptions.
 
-class DummyException(Exception):
-    pass
+Classes:
+    DummyError(Exception):
+        A custom exception class used as a default for exception handling.
 
-
-def import_global(
-        name, modules=None, exceptions=DummyException, locals_=None,
+Functions:
+    import_global(name, modules=None, exceptions=DummyError, locals_=None,
         globals_=None, level=-1):
-    '''Import the requested items into the global scope
+        Imports the requested items into the global scope, with support for
+        relative imports and custom exception handling.
+"""
+
+import typing
+
+from python_utils import _aliases
+
+
+class DummyError(Exception):
+    """A custom exception class used as a default for exception handling."""
+
+
+#: Backwards-compatible legacy alias for ``DummyError``.
+DummyException = DummyError
+
+
+def import_global(  # noqa: C901
+    name: str,
+    modules: list[str] | None = None,
+    exceptions: _aliases.ExceptionsType = DummyError,
+    locals_: _aliases.OptionalScope = None,
+    globals_: _aliases.OptionalScope = None,
+    level: int = -1,
+) -> typing.Any:  # sourcery skip: hoist-if-from-if
+    """Import the requested items into the global scope.
 
     WARNING! this method _will_ overwrite your global scope
-    If you have a variable named "path" and you call import_global('sys')
-    it will be overwritten with sys.path
+    If you have a variable named `path` and you call `import_global('sys')`
+    it will be overwritten with `sys.path`
 
     Args:
         name (str): the name of the module to import, e.g. sys
         modules (str): the modules to import, use None for everything
-        exception (Exception): the exception to catch, e.g. ImportError
-        `locals_`: the `locals()` method (in case you need a different scope)
-        `globals_`: the `globals()` method (in case you need a different scope)
+        exceptions (Exception): the exception to catch, e.g. ImportError
+        locals_: the `locals()` method (in case you need a different scope)
+        globals_: the `globals()` method (in case you need a different scope)
         level (int): the level to import from, this can be used for
         relative imports
-    '''
+    """
     frame = None
+    name_parts: list[str] = name.split('.')
+    modules_set: set[str] = set()
     try:
         # If locals_ or globals_ are not given, autodetect them by inspecting
         # the current stack
         if locals_ is None or globals_ is None:
             import inspect
+
             frame = inspect.stack()[1][0]
 
             if locals_ is None:
@@ -36,44 +66,52 @@ def import_global(
                 globals_ = frame.f_globals
 
         try:
-            name = name.split('.')
-
             # Relative imports are supported (from .spam import eggs)
-            if not name[0]:
-                name = name[1:]
+            if not name_parts[0]:
+                name_parts = name_parts[1:]
                 level = 1
 
             # raise IOError((name, level))
             module = __import__(
-                name=name[0] or '.',
+                name=name_parts[0] or '.',
                 globals=globals_,
                 locals=locals_,
-                fromlist=name[1:],
+                fromlist=name_parts[1:],
                 level=max(level, 0),
             )
 
             # Make sure we get the right part of a dotted import (i.e.
             # spam.eggs should return eggs, not spam)
             try:
-                for attr in name[1:]:
+                for attr in name_parts[1:]:
                     module = getattr(module, attr)
-            except AttributeError:
-                raise ImportError('No module named ' + '.'.join(name))
+            except AttributeError as e:
+                raise ImportError(
+                    'No module named ' + '.'.join(name_parts)
+                ) from e
 
             # If no list of modules is given, autodetect from either __all__
             # or a dir() of the module
             if not modules:
-                modules = getattr(module, '__all__', dir(module))
+                modules_set = set(getattr(module, '__all__', dir(module)))
             else:
-                modules = set(modules).intersection(dir(module))
+                modules_set = set(modules).intersection(dir(module))
 
             # Add all items in modules to the global scope
-            for k in set(dir(module)).intersection(modules):
+            for k in set(dir(module)).intersection(modules_set):
                 if k and k[0] != '_':
                     globals_[k] = getattr(module, k)
         except exceptions as e:
             return e
     finally:
         # Clean up, just to be sure
-        del name, modules, exceptions, locals_, globals_, frame
-
+        del (  # pyrefly: ignore[unsupported-delete]
+            name,
+            name_parts,
+            modules,
+            modules_set,
+            exceptions,
+            locals_,
+            globals_,
+            frame,
+        )