Completed
Pull Request — master (#272)
by Ryan
01:14
created

set_agg_backend()   A

Complexity

Conditions 1

Size

Total Lines 10

Duplication

Lines 0
Ratio 0 %

Importance

Changes 1
Bugs 0 Features 0
Metric Value
cc 1
dl 0
loc 10
rs 9.4285
c 1
b 0
f 0
1
# Copyright (c) 2008-2015 MetPy Developers.
2
# Distributed under the terms of the BSD 3-Clause License.
3
# SPDX-License-Identifier: BSD-3-Clause
4
r"""Collection of utilities for testing.
5
6
This includes:
7
* unit-aware test functions
8
* code for testing matplotlib figures
9
"""
10
11
import numpy.testing
12
from pint import DimensionalityError
13
import pytest
14
15
from .units import units
16
17
18
def check_and_drop_units(actual, desired):
19
    r"""Check that the units on the passed in arrays are compatible; return the magnitudes.
20
21
    Parameters
22
    ----------
23
    actual : `pint.Quantity` or array-like
24
25
    desired : `pint.Quantity` or array-like
26
27
    Returns
28
    -------
29
    actual, desired
30
        array-like versions of `actual` and `desired` once they have been
31
        coerced to compatible units.
32
33
    Raises
34
    ------
35
    AssertionError
36
        If the units on the passed in objects are not compatible.
37
    """
38
    try:
39
        # If the desired result has units, add dimensionless units if necessary, then
40
        # ensure that this is compatible to the desired result.
41
        if hasattr(desired, 'units'):
42
            if not hasattr(actual, 'units'):
43
                actual = units.Quantity(actual, 'dimensionless')
44
            actual = actual.to(desired.units)
45
        # Otherwise, the desired result has no units. Convert the actual result to
46
        # dimensionless units if it is a united quantity.
47
        else:
48
            if hasattr(actual, 'units'):
49
                actual = actual.to('dimensionless')
50
    except DimensionalityError:
51
        raise AssertionError('Units are not compatible: {} should be {}'.format(actual.units,
52
                                                                                desired.units))
53
    except AttributeError:
54
        pass
55
56
    if hasattr(actual, 'magnitude'):
57
        actual = actual.magnitude
58
    if hasattr(desired, 'magnitude'):
59
        desired = desired.magnitude
60
61
    return actual, desired
62
63
64
def assert_almost_equal(actual, desired, decimal=7):
65
    """Check that values are almost equal, including units.
66
67
    Wrapper around :func:`numpy.testing.assert_almost_equal`
68
    """
69
    actual, desired = check_and_drop_units(actual, desired)
70
    numpy.testing.assert_almost_equal(actual, desired, decimal)
71
72
73
def assert_array_almost_equal(actual, desired, decimal=7):
74
    """Check that arrays are almost equal, including units.
75
76
    Wrapper around :func:`numpy.testing.assert_array_almost_equal`
77
    """
78
    actual, desired = check_and_drop_units(actual, desired)
79
    numpy.testing.assert_array_almost_equal(actual, desired, decimal)
80
81
82
def assert_array_equal(actual, desired):
83
    """Check that arrays are equal, including units.
84
85
    Wrapper around :func:`numpy.testing.assert_array_equal`
86
    """
87
    actual, desired = check_and_drop_units(actual, desired)
88
    numpy.testing.assert_array_equal(actual, desired)
89
90
91
@pytest.fixture(scope='module', autouse=True)
92
def set_agg_backend():
93
    """Fixture to ensure the Agg backend is active."""
94
    import matplotlib.pyplot as plt
95
    prev_backend = plt.get_backend()
96
    try:
97
        plt.switch_backend('agg')
98
        yield
99
    finally:
100
        plt.switch_backend(prev_backend)
101