_config.py

"""Global configuration state and functions for management
"""
import os
from contextlib import contextmanager as contextmanager

_global_config = {
    'assume_finite': bool(os.environ.get('SKLEARN_ASSUME_FINITE', False)),
}


def get_config():
    """Retrieve current values for configuration set by :func:`set_config`

    Returns
    -------
    config : dict
        Keys are parameter names that can be passed to :func:`set_config`.
    """
    return _global_config.copy()


def set_config(assume_finite=None):
    """Set global scikit-learn configuration

    Parameters
    ----------
    assume_finite : bool, optional
        If True, validation for finiteness will be skipped,
        saving time, but leading to potential crashes. If
        False, validation for finiteness will be performed,
        avoiding error.  Global default: False.
    """
    if assume_finite is not None:
        _global_config['assume_finite'] = assume_finite


@contextmanager
def config_context(**new_config):
    """Context manager for global scikit-learn configuration

    Parameters
    ----------
    assume_finite : bool, optional
        If True, validation for finiteness will be skipped,
        saving time, but leading to potential crashes. If
        False, validation for finiteness will be performed,
        avoiding error.  Global default: False.

    Notes
    -----
    All settings, not just those presently modified, will be returned to
    their previous values when the context manager is exited. This is not
    thread-safe.

    Examples
    --------
    >>> import sklearn
    >>> from sklearn.utils.validation import assert_all_finite
    >>> with sklearn.config_context(assume_finite=True):
    ...     assert_all_finite([float('nan')])
    >>> with sklearn.config_context(assume_finite=True):
    ...     with sklearn.config_context(assume_finite=False):
    ...         assert_all_finite([float('nan')])
    ... # doctest: +ELLIPSIS
    Traceback (most recent call last):
    ...
    ValueError: Input contains NaN, ...
    """
    old_config = get_config().copy()
    set_config(**new_config)

    try:
        yield
    finally:
        set_config(**old_config)