-
Notifications
You must be signed in to change notification settings - Fork 258
/
tst_enum.py
67 lines (60 loc) · 2.31 KB
/
tst_enum.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
import sys
import unittest
import os
import tempfile
from netCDF4 import Dataset
import numpy as np
from numpy.testing import assert_array_equal
FILE_NAME = tempfile.NamedTemporaryFile(suffix='.nc', delete=False).name
ENUM_NAME = 'cloud_t'
ENUM_BASETYPE = np.int8
VAR_NAME = 'primary_cloud'
ENUM_DICT = {u'Altocumulus': 7, u'Missing': 127, u'Stratus': 2, u'Clear': 0,
u'Nimbostratus': 6, u'Cumulus': 4, u'Altostratus': 5, u'Cumulonimbus': 1,
u'Stratocumulus': 3}
datain = np.array([ENUM_DICT['Clear'],ENUM_DICT['Stratus'],ENUM_DICT['Cumulus'],\
ENUM_DICT['Missing'],ENUM_DICT['Cumulonimbus']],dtype=ENUM_BASETYPE)
datain_masked = np.ma.masked_values(datain,ENUM_DICT['Missing'])
class EnumTestCase(unittest.TestCase):
def setUp(self):
self.file = FILE_NAME
f = Dataset(self.file,'w')
cloud_type = f.createEnumType(ENUM_BASETYPE,ENUM_NAME,ENUM_DICT)
# make sure KeyError raised if non-integer basetype used.
try:
cloud_typ2 = f.createEnumType(np.float32,ENUM_NAME,ENUM_DICT)
except KeyError:
pass
f.createDimension('time',None)
cloud_var =\
f.createVariable(VAR_NAME,cloud_type,'time',\
fill_value=ENUM_DICT['Missing'])
cloud_var[:] = datain_masked
# make sure ValueError raised if illegal value assigned to Enum var.
try:
cloud_var[cloud_var.shape[0]] = 99
except ValueError:
pass
f.close()
def tearDown(self):
# Remove the temporary files
os.remove(self.file)
def runTest(self):
"""testing enum data type"""
f = Dataset(self.file, 'r')
v = f.variables[VAR_NAME]
assert v.datatype.enum_dict == ENUM_DICT
assert list(f.enumtypes.keys()) == [ENUM_NAME]
assert f.enumtypes[ENUM_NAME].name == ENUM_NAME # issue 775
assert f.enumtypes[ENUM_NAME].dtype == ENUM_BASETYPE
assert v._FillValue == ENUM_DICT['Missing']
v.set_auto_mask(False)
data = v[:]
assert_array_equal(data, datain)
v.set_auto_mask(True) # check to see if auto masking works
data = v[:]
assert_array_equal(data, datain_masked)
assert_array_equal(data.mask, datain_masked.mask)
f.close()
if __name__ == '__main__':
unittest.main()