<commit_before>import os
import warnings

import numpy as np
from numpy.testing.noseclasses import KnownFailureTest

import scipy.special as sc

#------------------------------------------------------------------------------
# Enable convergence and loss of precision warnings -- turn off one by one
#------------------------------------------------------------------------------

def with_special_errors(func):
    """
    Enable special function errors (such as underflow, overflow,
    loss of precision, etc.)
    """
    def wrapper(*a, **kw):
        old_filters = list(getattr(warnings, 'filters', []))
        old_errprint = sc.errprint(1)
        warnings.filterwarnings("error", category=sc.SpecialFunctionWarning)
        try:
            return func(*a, **kw)
        finally:
            sc.errprint(old_errprint)
            setattr(warnings, 'filters', old_filters)
    wrapper.__name__ = func.__name__
    wrapper.__doc__ = func.__doc__
    return wrapper

#------------------------------------------------------------------------------
# Comparing function values at many data points at once, with helpful
#------------------------------------------------------------------------------

def assert_tol_equal(a, b, rtol=1e-7, atol=0, err_msg='', verbose=True):
    """Assert that `a` and `b` are equal to tolerance ``atol + rtol*abs(b)``"""
    def compare(x, y):
        return np.allclose(x, y, rtol=rtol, atol=atol)
    a, b = np.asanyarray(a), np.asanyarray(b)
    header = 'Not equal to tolerance rtol=%g, atol=%g' % (rtol, atol)
    np.testing.utils.assert_array_compare(compare, a, b, err_msg=str(err_msg),
                                          verbose=verbose, header=header)

#------------------------------------------------------------------------------
# Comparing function values at many data points at once, with helpful
# error reports
#------------------------------------------------------------------------------<commit_msg>ENH: special: add assert_func_equal to clean up the way dataset tests are implemented<commit_after>import os
import warnings

import numpy as np
from numpy.testing.noseclasses import KnownFailureTest

import scipy.special as sc

__all__ = ['with_special_errors', 'assert_tol_equal', 'assert_func_equal',
           'FuncData']

#------------------------------------------------------------------------------
# Enable convergence and loss of precision warnings -- turn off one by one
#------------------------------------------------------------------------------

def with_special_errors(func):
    """
    Enable special function errors (such as underflow, overflow,
    loss of precision, etc.)
    """
    def wrapper(*a, **kw):
        old_filters = list(getattr(warnings, 'filters', []))
        old_errprint = sc.errprint(1)
        warnings.filterwarnings("error", category=sc.SpecialFunctionWarning)
        try:
            return func(*a, **kw)
        finally:
            sc.errprint(old_errprint)
            setattr(warnings, 'filters', old_filters)
    wrapper.__name__ = func.__name__
    wrapper.__doc__ = func.__doc__
    return wrapper

#------------------------------------------------------------------------------
# Comparing function values at many data points at once, with helpful
#------------------------------------------------------------------------------

def assert_tol_equal(a, b, rtol=1e-7, atol=0, err_msg='', verbose=True):
    """Assert that `a` and `b` are equal to tolerance ``atol + rtol*abs(b)``"""
    def compare(x, y):
        return np.allclose(x, y, rtol=rtol, atol=atol)
    a, b = np.asanyarray(a), np.asanyarray(b)
    header = 'Not equal to tolerance rtol=%g, atol=%g' % (rtol, atol)
    np.testing.utils.assert_array_compare(compare, a, b, err_msg=str(err_msg),
                                          verbose=verbose, header=header)

#------------------------------------------------------------------------------
# Comparing function values at many data points at once, with helpful
# error reports
#------------------------------------------------------------------------------

def assert_func_equal(func, results, points, rtol=None, atol=None,
                      param_filter=None, knownfailure=None,
                      vectorized=True, dtype=None):
    if hasattr(points, 'next'):
        # it's a generator
        points = list(points)

    points = np.asarray(points)
    if points.ndim == 1:
        points = points[:,None]

    if hasattr(results, '__name__'):
        # function
        if vectorized:
            results = results(*tuple(points.T))
        else:
            results = np.array([results(*tuple(p)) for p in points])
            if results.dtype == object:
                try:
                    results = results.astype(float)
                except TypeError:
                    results = results.astype(complex)
    else:
        results = np.asarray(results)

    npoints = points.shape[1]

    data = np.c_[points, results]
    fdata = FuncData(func, data, range(npoints), range(npoints, data.shape[1]),
                     rtol=rtol, atol=atol, param_filter=param_filter,
                     knownfailure=knownfailure)
    fdata.check()
{'fraction_non_alphanumeric': 0.27155504234026173, 'fraction_numerical': 0.002117013086989992, 'mean_word_length': 1.4985576923076922, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4764359', 'n_tokens_mistral': 1316, 'n_tokens_neox': 1295, 'n_words': 358}