1# Copyright (c) 2015,2016,2017,2019 MetPy Developers.
2# Distributed under the terms of the BSD 3-Clause License.
3# SPDX-License-Identifier: BSD-3-Clause
4"""Tests for the `ctables` module."""
5
6from io import StringIO
7from pathlib import Path
8import tempfile
9
10import numpy as np
11import pytest
12
13from metpy.plots.ctables import ColortableRegistry, convert_gempak_table
14
15
16@pytest.fixture()
17def registry():
18    """Set up a registry for use by the tests."""
19    return ColortableRegistry()
20
21
22def test_package_resource(registry):
23    """Test registry scanning package resource."""
24    registry.scan_resource('metpy.plots', 'nexrad_tables')
25    assert 'cc_table' in registry
26
27
28def test_scan_dir(registry):
29    """Test registry scanning a directory and ignoring files it can't handle ."""
30    try:
31        kwargs = {'mode': 'w', 'dir': '.', 'suffix': '.tbl', 'delete': False, 'buffering': 1}
32        with tempfile.NamedTemporaryFile(**kwargs) as fobj:
33            fobj.write('"red"\n"lime"\n"blue"\n')
34            good_file = Path(fobj.name)
35
36        # Unrelated table file that *should not* impact the scan
37        with tempfile.NamedTemporaryFile(**kwargs) as fobj:
38            fobj.write('PADK     704540 ADAK NAS\n')
39            bad_file = Path(fobj.name)
40
41        # Needs to be outside with so it's closed on windows
42        registry.scan_dir(good_file.parent)
43        name = good_file.with_suffix('').name
44
45        assert name in registry
46        assert registry[name] == [(1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, 0.0, 1.0)]
47    finally:
48        good_file.unlink()
49        bad_file.unlink()
50
51
52def test_read_file(registry):
53    """Test reading a colortable from a file."""
54    fobj = StringIO('(0., 0., 1.0)\n"red"\n"#0000FF" #Blue')
55
56    registry.add_colortable(fobj, 'test_table')
57
58    assert 'test_table' in registry
59    assert registry['test_table'] == [(0.0, 0.0, 1.0), (1.0, 0.0, 0.0), (0.0, 0.0, 1.0)]
60
61
62def test_read_bad_file(registry):
63    """Test what error results when reading a malformed file."""
64    with pytest.raises(RuntimeError):
65        fobj = StringIO('PADK     704540 ADAK NAS                         '
66                        'AK US  5188 -17665     4  0')
67        registry.add_colortable(fobj, 'sfstns')
68
69
70def test_get_colortable(registry):
71    """Test getting a colortable from the registry."""
72    true_colors = [(0.0, 0.0, 1.0), (1.0, 0.0, 0.0)]
73    registry['table'] = true_colors
74
75    table = registry.get_colortable('table')
76    assert table.N == 2
77    assert table.colors == true_colors
78
79
80def test_get_steps(registry):
81    """Test getting a colortable and norm with appropriate steps."""
82    registry['table'] = [(0.0, 0.0, 1.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)]
83    norm, cmap = registry.get_with_steps('table', 5., 10.)
84    assert cmap(norm(np.array([6.]))).tolist() == [[0.0, 0.0, 1.0, 1.0]]
85    assert cmap(norm(np.array([14.9]))).tolist() == [[0.0, 0.0, 1.0, 1.0]]
86    assert cmap(norm(np.array([15.1]))).tolist() == [[1.0, 0.0, 0.0, 1.0]]
87    assert cmap(norm(np.array([26.]))).tolist() == [[0.0, 1.0, 0.0, 1.0]]
88
89
90def test_get_steps_negative_start(registry):
91    """Test bad start for get with steps (issue #81)."""
92    registry['table'] = [(0.0, 0.0, 1.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)]
93    norm, _ = registry.get_with_steps('table', -10, 5)
94    assert norm.vmin == -10
95    assert norm.vmax == 5
96
97
98def test_get_range(registry):
99    """Test getting a colortable and norm with appropriate range."""
100    registry['table'] = [(0.0, 0.0, 1.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)]
101    norm, cmap = registry.get_with_range('table', 5., 35.)
102    assert cmap(norm(np.array([6.]))).tolist() == [[0.0, 0.0, 1.0, 1.0]]
103    assert cmap(norm(np.array([14.9]))).tolist() == [[0.0, 0.0, 1.0, 1.0]]
104    assert cmap(norm(np.array([15.1]))).tolist() == [[1.0, 0.0, 0.0, 1.0]]
105    assert cmap(norm(np.array([26.]))).tolist() == [[0.0, 1.0, 0.0, 1.0]]
106
107
108def test_get_boundaries(registry):
109    """Test getting a colortable with explicit boundaries."""
110    registry['table'] = [(0.0, 0.0, 1.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)]
111    norm, cmap = registry.get_with_boundaries('table', [0., 8., 10., 20.])
112    assert cmap(norm(np.array([7.]))).tolist() == [[0.0, 0.0, 1.0, 1.0]]
113    assert cmap(norm(np.array([9.]))).tolist() == [[1.0, 0.0, 0.0, 1.0]]
114    assert cmap(norm(np.array([10.1]))).tolist() == [[0.0, 1.0, 0.0, 1.0]]
115
116
117def test_gempak():
118    """Test GEMPAK colortable conversion."""
119    infile = StringIO("""!   wvcolor.tbl
120                         0      0      0
121                       255    255    255
122                       """)
123    outfile = StringIO()
124
125    # Do the conversion
126    convert_gempak_table(infile, outfile)
127
128    # Reset and grab contents
129    outfile.seek(0)
130    result = outfile.read()
131
132    assert result == '(0.000000, 0.000000, 0.000000)\n(1.000000, 1.000000, 1.000000)\n'
133