import difflib

import functools

import sys

import numbers

import numpy as np

from .misc import indent

__all__ = ['fixed_width_indent', 'diff_values', 'report_diff_values',

'where_not_allclose']

# Smaller default shift-width for indent

fixed_width_indent = functools.partial(indent, width=2)

def diff_values(a, b, rtol=0.0, atol=0.0):

"""

Diff two scalar values. If both values are floats, they are compared to

within the given absolute and relative tolerance.

Parameters

----------

a, b : int, float, str

Scalar values to compare.

rtol, atol : float

Relative and absolute tolerances as accepted by

:func:`numpy.allclose`.

Returns

-------

is_different : bool

`True` if they are different, else `False`.

"""

if isinstance(a, float) and isinstance(b, float):

if np.isnan(a) and np.isnan(b):

return False

return not np.allclose(a, b, rtol=rtol, atol=atol)

else:

return a != b

def report_diff_values(a, b, fileobj=sys.stdout, indent_width=0):

"""

Write a diff report between two values to the specified file-like object.

Parameters

----------

a, b

Values to compare. Anything that can be turned into strings

and compared using :py:mod:`difflib` should work.

fileobj : obj

File-like object to write to.

The default is ``sys.stdout``, which writes to terminal.

indent_width : int

Character column(s) to indent.

Returns

-------

identical : bool

`True` if no diff, else `False`.

"""

if isinstance(a, np.ndarray) and isinstance(b, np.ndarray):

if a.shape != b.shape:

fileobj.write(

fixed_width_indent(' Different array shapes:\n',

indent_width))

report_diff_values(str(a.shape), str(b.shape), fileobj=fileobj,

indent_width=indent_width + 1)

return False

diff_indices = np.transpose(np.where(a != b))

num_diffs = diff_indices.shape[0]

for idx in diff_indices[:3]:

lidx = idx.tolist()

fileobj.write(

fixed_width_indent(' at {!r}:\n'.format(lidx), indent_width))

report_diff_values(a[tuple(idx)], b[tuple(idx)], fileobj=fileobj,

indent_width=indent_width + 1)

if num_diffs > 3:

fileobj.write(fixed_width_indent(

' ...and at {:d} more indices.\n'.format(num_diffs - 3),

indent_width))

return False

return num_diffs == 0

typea = type(a)

typeb = type(b)

if typea == typeb:

lnpad = ' '

sign_a = 'a>'

sign_b = 'b>'

if isinstance(a, numbers.Number):

a = repr(a)

b = repr(b)

else:

a = str(a)

b = str(b)

else:

padding = max(len(typea.__name__), len(typeb.__name__)) + 3

lnpad = (padding + 1) * ' '

sign_a = ('(' + typea.__name__ + ') ').rjust(padding) + 'a>'

sign_b = ('(' + typeb.__name__ + ') ').rjust(padding) + 'b>'

is_a_str = isinstance(a, str)

is_b_str = isinstance(b, str)

a = (repr(a) if ((is_a_str and not is_b_str) or

(not is_a_str and isinstance(a, numbers.Number)))

else str(a))

b = (repr(b) if ((is_b_str and not is_a_str) or

(not is_b_str and isinstance(b, numbers.Number)))

else str(b))

identical = True

for line in difflib.ndiff(a.splitlines(), b.splitlines()):

if line[0] == '-':

identical = False

line = sign_a + line[1:]

elif line[0] == '+':

identical = False

line = sign_b + line[1:]

else:

line = lnpad + line

fileobj.write(fixed_width_indent(

' {}\n'.format(line.rstrip('\n')), indent_width))

return identical

def where_not_allclose(a, b, rtol=1e-5, atol=1e-8):

"""

A version of :func:`numpy.allclose` that returns the indices

where the two arrays differ, instead of just a boolean value.

Parameters

----------

a, b : array_like

Input arrays to compare.

rtol, atol : float

Relative and absolute tolerances as accepted by

:func:`numpy.allclose`.

Returns

-------

idx : tuple of arrays

Indices where the two arrays differ.

"""

# Create fixed mask arrays to handle INF and NaN; currently INF and NaN

# are handled as equivalent

if not np.all(np.isfinite(a)):

a = np.ma.fix_invalid(a).data

if not np.all(np.isfinite(b)):

b = np.ma.fix_invalid(b).data

if atol == 0.0 and rtol == 0.0:

# Use a faster comparison for the most simple (and common) case

return np.where(a != b)

return np.where(np.abs(a - b) > (atol + rtol * np.abs(b)))