Source code for

Miscellaneous python functions 

.. autosummary::


.. codeauthor:: David Zwicker <>

import errno
import functools
import importlib
import json
import os
import unittest
from pathlib import Path
from typing import Any, Callable, Dict, Optional, Sequence, TypeVar, Union

import numpy as np

from .typing import ArrayLike, Number

TFunc = TypeVar("TFunc", bound=Callable[..., Any])

[docs]def module_available(module_name: str) -> bool: """check whether a python module is available Args: module_name (str): The name of the module Returns: `True` if the module can be imported and `False` otherwise """ try: importlib.import_module(module_name) except ImportError: return False else: return True
[docs]def ensure_directory_exists(folder: Union[str, Path]): """creates a folder if it not already exists Args: folder (str): path of the new folder """ folder = str(folder) if folder == "": return try: os.makedirs(folder) except OSError as err: if err.errno != errno.EEXIST: raise
[docs]def preserve_scalars(method: TFunc) -> TFunc: """decorator that makes vectorized methods work with scalars This decorator allows to call functions that are written to work on numpy arrays to also accept python scalars, like `int` and `float`. Essentially, this wrapper turns them into an array and unboxes the result. Args: method: The method being decorated Returns: The decorated method """ @functools.wraps(method) def wrapper(self, *args): args = [number_array(arg, copy=False) for arg in args] if args[0].ndim == 0: args = [arg[None] for arg in args] return method(self, *args)[0] else: return method(self, *args) return wrapper # type: ignore
[docs]def decorator_arguments(decorator: Callable) -> Callable: r"""make a decorator usable with and without arguments: The resulting decorator can be used like `@decorator` or `@decorator(\*args, \**kwargs)` Inspired by Args: decorator: the decorator that needs to be modified Returns: the decorated function """ @functools.wraps(decorator) def new_decorator(*args, **kwargs): if len(args) == 1 and len(kwargs) == 0 and callable(args[0]): # actual decorated function return decorator(args[0]) else: # decorator arguments return lambda realf: decorator(realf, *args, **kwargs) return new_decorator
[docs]def skipUnlessModule( module_names: Union[Sequence[str], str] ) -> Callable[[TFunc], TFunc]: """decorator that skips a test when a module is not available Args: module_names (str): The name of the required module(s) Returns: A function, so this can be used as a decorator """ if isinstance(module_names, str): module_names = [module_names] for module_name in module_names: if not module_available(module_name): # return decorator skipping test return unittest.skip(f"requires {module_name}") # return no-op decorator if all modules are available def wrapper(f: TFunc) -> TFunc: return f return wrapper
[docs]def import_class(identifier: str): """import a class or module given an identifier Args: identifier (str): The identifier can be a module or a class. For instance, calling the function with the string `identifier == 'numpy.linalg.norm'` is roughly equivalent to running `from numpy.linalg import norm` and would return a reference to `norm`. """ module_path, _, class_name = identifier.rpartition(".") if module_path: module = importlib.import_module(module_path) return getattr(module, class_name) else: # this happens when identifier does not contain a dot return importlib.import_module(class_name)
[docs]class classproperty(property): """decorator that can be used to define read-only properties for classes. This is inspired by the implementation of :mod:`astropy`, see ` <>`_. Example: The decorator can be used much like the `property` decorator:: class Test(): item: str = 'World' @classproperty def message(cls): return 'Hello ' + cls.item print(Test.message) """ def __new__(cls, fget=None, doc=None): if fget is None: # use wrapper to support decorator without arguments def wrapper(func): return cls(func) return wrapper return super().__new__(cls) def __init__(self, fget, doc=None): fget = self._wrap_fget(fget) super().__init__(fget=fget, doc=doc) if doc is not None: self.__doc__ = doc def __get__(self, obj, objtype): # The base property.__get__ will just return self here; # instead we pass objtype through to the original wrapped # function (which takes the class as its sole argument) return self.fget.__wrapped__(objtype)
[docs] def getter(self, fget): return super().getter(self._wrap_fget(fget))
[docs] def setter(self, fset): raise NotImplementedError("classproperty is read-only")
[docs] def deleter(self, fdel): raise NotImplementedError("classproperty is read-only")
@staticmethod def _wrap_fget(orig_fget): if isinstance(orig_fget, classmethod): orig_fget = orig_fget.__func__ @functools.wraps(orig_fget) def fget(obj): return orig_fget(obj.__class__) return fget
[docs]class hybridmethod: """ descriptor that can be used as a decorator to allow calling a method both as a classmethod and an instance method Adapted from """ def __init__(self, fclass, finstance=None, doc=None): self.fclass = fclass self.finstance = finstance self.__doc__ = doc or fclass.__doc__ # support use on abstract base classes self.__isabstractmethod__ = bool(getattr(fclass, "__isabstractmethod__", False)) # # support extracting information about the function for sphinx if finstance: self.__func__ = finstance
[docs] def classmethod(self, fclass): return type(self)(fclass, self.finstance, None)
[docs] def instancemethod(self, finstance): return type(self)(self.fclass, finstance, self.__doc__)
def __get__(self, instance, cls): if instance is None or self.finstance is None: # either bound to the class, or no instance method available return self.fclass.__get__(cls, None) return self.finstance.__get__(instance, cls)
[docs]def estimate_computation_speed(func: Callable, *args, **kwargs) -> float: """estimates the computation speed of a function Args: func (callable): The function to call Returns: float: the number of times the function can be calculated in one second. The inverse is thus the runtime in seconds per function call """ import timeit test_duration = kwargs.pop("test_duration", 1) # prepare the function if args or kwargs: test_func = functools.partial(func, *args, **kwargs) else: test_func = func # type: ignore # call function once to allow caches be filled test_func() # call the function until the total time is achieved number, duration = 1, 0 while duration < 0.1 * test_duration: number *= 10 duration = timeit.timeit(test_func, number=number) # type: ignore return number / duration
[docs]def hdf_write_attributes( hdf_path, attributes: Optional[Dict[str, Any]] = None, raise_serialization_error: bool = False, ) -> None: """write (JSON-serialized) attributes to a hdf file Args: hdf_path: Path to a group or dataset in an open HDF file attributes (dict): Dictionary with values written as attributes raise_serialization_error (bool): Flag indicating whether serialization errors are raised or silently ignored """ if attributes is None: return for key, value in attributes.items(): try: value_serialized = json.dumps(value) except TypeError: if raise_serialization_error: raise else: hdf_path.attrs[key] = value_serialized
[docs]def number(value: Union[Number, str]) -> Number: """convert a value into a float or complex number Args: value (Number or str): The value which needs to be converted Result: Number: A complex number or a float if the imaginary part vanishes """ result = complex(value) return result.real if result.imag == 0 else result
[docs]def get_common_dtype(*args): r"""returns a dtype in which all arguments can be represented Args: *args: All items (arrays, scalars, etc) to be checked Returns: numpy.cdouble if any entry is complex, otherwise np.double """ for arg in args: if np.iscomplexobj(arg): return np.cdouble return np.double
[docs]def number_array(data: ArrayLike, dtype=None, copy: bool = True) -> np.ndarray: """convert an array with arbitrary dtype either to np.double or np.cdouble Args: data (:class:`~numpy.ndarray`): The data that needs to be converted to a float array. This can also be any iterable of numbers. dtype (numpy dtype): The data type of the field. All the numpy dtypes are supported. If omitted, it will be determined from `data` automatically. copy (bool): Whether the data must be copied (in which case the original array is left untouched). Note that data will always be copied when changing the dtype. Returns: :class:`~numpy.ndarray`: An array with the correct dtype """ if dtype is None: # dtype needs to be determined automatically try: # convert the result to a numpy array with the given dtype result = np.array(data, dtype=get_common_dtype(data), copy=copy) except TypeError: # Conversion can fail when `data` contains a complex sympy number, i.e., # sympy.I. In this case, we simply try to convert the expression using a # complex dtype result = np.array(data, dtype=np.cdouble, copy=copy) else: # a specific dtype is requested result = np.array(data, dtype=np.dtype(dtype), copy=copy) return result