tools.py

# coding: utf-8
"""Universal toolkit.

The functions in this module are general helper functions, which may be useful
in other projects as well.
"""

import sys
import re
import time
from math import log, floor
from contextlib import contextmanager


def user(message, dtype=str, allow_interrupt=False):
    """User input with type check.

    Parameters
    ----------
    message : str
        message to be displayed to the user
    dtype : callable, optional
        Any callable type the input value should be cast to. If the expected
        type doesn't match, `dtype` should rise a ValueError.
        Defaults to str
    allow_interrupt : bool, optional
        If set, the user may cancel the query safely with Keyboard Interrupt or
        EOF

    Raises
    ------
    KeyboardError, EOFError
        if interrupts are not allowed and the user cancels input using
        Ctrl+C (Keyboard Interrupt) or Ctrl+D (End of File character)

    Returns
    -------
    value : `dtype`
        if no interrupts are allowed
    success, value : (bool, `dtype`)
        if interrupts are allowed

    Examples
    --------
    >>> num = user("Your favourite number, please", dtype=int)
    Your favourite number, please: fourty-two
    <class 'int'> expected, try again.
    Your favourite number, please: 42
    >>> print(num)
    42
    """
    while True:
        try:
            x = input(message + ": ")
            try:
                val = dtype(x)
            except ValueError:
                print("{!r} expected, try again.\n".format(dtype))
                continue
            return (True, val) if allow_interrupt else val
        except (KeyboardInterrupt, EOFError):
            if allow_interrupt:
                return False, None
            raise

def hr(value, fmt='.0f', sep=' ', base=1000):
    """Format a value into human-readable form.

    From a list of SI-prefixes (kilo, Mega, Giga, milli, micro, nano, etc.), the
    matching one is being taken and appended to the accordingly scaled value.

    Parameters
    ----------
    value : int or float
        value to process
    fmt : str, optional
        format of the value's string representation, defaults to '.0f'
    sep : str, optional
        separator between value and unit, defaults to space (' ')
    base : int, optional
        base for 'kilo'-prefix, e.g. 1000 or 1024. defaults to 1000

    Returns
    -------
    hr_value : str
        value formatted to a human-friendly form

    Examples
    --------
    >>> hr(17.5e-3) + 'm'
    '18 mm'
    >>> hr(17.5e-3, fmt='.3f') + 'm'
    '17.500 mm'
    >>> hr(-13.57e3, fmt='.1f') + 'N'
    '-13.6 kN'
    >>> hr(4118396, fmt='.2f', base=1024) + 'iB'
    '3.93 MiB'
    """
    mag = floor(log(abs(value), base))
    if mag > 6:
        mag = 6
    if mag < -6:
        mag = -6

    # mag      0   1    2    3    4    5    6   -6   -5   -4   -3   -2   -1
    prefix = ['', 'k', 'M', 'G', 'T', 'P', 'E', 'a', 'f', 'p', 'n', 'µ', 'm']

    fmt_string = '{{0:{}}}{}{}'.format(fmt, sep, prefix[mag])
    return fmt_string.format(value / base**mag)

@contextmanager
def msg(message, end="done.", timing=True):
    """Informative timer as context manager.

    Parameters
    ----------
    message : str
        message to be displayed
    end : str, optional
        message to be added upon completion, defaults to 'done.'
    timing : bool, optional
        stop the execution time and display afterwards

    Examples
    --------
    >>> with msg("counting"):
    ...     for i in range(int(1e6)):
    ...         pass
    ...
    counting... done.                                                 [ 93.7 µs]

    """
    print(message, end='... ', flush=True)
    t_start = time.time()
    yield
    if timing:
        t = time.time() - t_start
        end = "{s:{d}s}[{hrt:<7s}s]".format(s=end, d=80-len(message)-15,
                                            hrt=hr(t, fmt='5.1f'))
    print(end)

def get(dictionary, *keys, default=None):
    """A `dict.get()` with multiple keys to check before falling back to the
    default value

    Parameters
    ----------
    dictionary : dict
        dictionary to extract data from
    *keys : sequence of hashable types
        sequence of keys to check
    default : any type, optional
        value to return if none of the `keys` was found in the dictionary,
        defaults to None.

    Returns
    -------
    value :
        dictionary item of the first matching key, or `default` if no key
        matches
    """
    # default = keys.pop()
    for key in keys:
        if key in dictionary:
            return dictionary[key]
    return default

def range_definition_to_list(range_definition):
    """Convert a range definition into a list of integers.

    A range definition consists of comma-separated values and ranges.
    See Examples section for details.

    Parameters
    ----------
    range_definition : str
        textual definition of range

    Returns
    -------
    range_list : list of integers
        specified range

    Raises
    ------
    ValueError
        if (parts of) the range definition were invalid

    Examples
    --------
    >>> range_definition_to_list('1,3-5,7')
    [1, 3, 4, 5, 7]
    >>> range_definition_to_list('1,3-5+7')
    ValueError: '1,3-5+7' is not a valid range definition
    """
    range_list = []
    for value in range_definition.split(','):
        value = value.strip()
        if value.isdigit():
            range_list.append(int(value))
        elif re.match(r'^\d+\s*\-\s*\d+$', value):
            start, end = (int(val) for val in value.split('-'))
            range_list += list(range(start, end+1))
        else:
            raise ValueError("'{}' is not a valid range definition" \
                    .format(range_definition))
    return range_list

def zip_longest_repeat_last(*iterables):
    """zip a sequence of iterables, repeating last values until the largest
    iterable is exhausted.

    Parameters
    ----------
    *iterables : sequence
        sequence of iterable types

    Yields
    ------
    items : tuple
        zipped items of `iterables`

    Notes
    -----
    `len()` is being applied on every iterable to determine total steps.

    Examples
    --------
    >>> list(zip_longes_repeat_last(range(4), 'abc'))
    [(0, 'a'), (1, 'b'), (2, 'c'), (3, 'c')]
    """
    count = max(len(it) for it in iterables)
    def repeater(iterable):
        """Iterate through `iterable`, then repeat the last item"""
        yield from iterable
        for _ in range(count - len(iterable)):
            yield iterable[-1]
    yield from zip(*(repeater(it) for it in iterables))

def as_bool(value):
    """Parse human-entered boolean values.

    Evaluates numeric values as well as 'yes', 'y' and 'true'.
    """
    if isinstance(value, str):
        try:
            return bool(float(value))
        except ValueError:
            return value.lower() in ('y', 'yes', 'true', 'j', 'ja')
    return bool(value)

def docstring_parameter(*sub, **kwsub):
    """Decorator for dynamic docstring generation.

    Wraps any object to format its docstring using `.format(*args, **kwargs)`.
    Seen at http://stackoverflow.com/a/10308363/1525423
    """
    def decorator(obj):
        """Rewrites a `obj`'s docstring"""
        obj.__doc__ = obj.__doc__.format(*sub, **kwsub)
        return obj
    return decorator

def exit_why(why, code=1):
    """Display a message and exit with exit code."""
    print(why)
    sys.exit(code)