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