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
VL_NAME = 'vlen_type'
VL_BASETYPE = np.int16
DIM1_NAME = 'lon'
DIM2_NAME = 'lat'
nlons = 5; nlats = 5
VAR1_NAME = 'ragged'
VAR2_NAME = 'strings'
VAR3_NAME = 'strings_alt'
VAR4_NAME = 'string_scalar'
VAR5_NAME = 'vlen_scalar'
data = np.empty(nlats*nlons,object)
datas = np.empty(nlats*nlons,object)
nn = 0
for n in range(nlats*nlons):
    nn = nn + 1
    data[n] = np.arange(nn,dtype=VL_BASETYPE)
    datas[n] = ''.join([chr(i) for i in range(97,97+nn+1)])
data = np.reshape(data,(nlats,nlons))
datas = np.reshape(datas,(nlats,nlons))


class VariablesTestCase(unittest.TestCase):

    def setUp(self):
        self.file = FILE_NAME
        f = Dataset(self.file,'w')
        vlen_t = f.createVLType(VL_BASETYPE, VL_NAME)
        f.createDimension(DIM1_NAME,nlons)
        f.createDimension(DIM2_NAME,nlats)
        ragged = f.createVariable(VAR1_NAME, vlen_t,\
                (DIM2_NAME,DIM1_NAME))
        strings = f.createVariable(VAR2_NAME, str,
                (DIM2_NAME,DIM1_NAME))
        strings_alt = f.createVariable(VAR3_NAME, datas.astype(str).dtype,
                                       (DIM2_NAME, DIM1_NAME))
        string_scalar = f.createVariable(VAR4_NAME,str,())
        vlen_scalar = f.createVariable(VAR5_NAME,vlen_t,())
        ragged[:] = data
        ragged[-1,-1] = data[-1,-1]
        strings[:] = datas
        strings[-2,-2] = datas[-2,-2]
        strings_alt[:] = datas.astype(str)
        string_scalar[...] = 'foo'  #issue458
        vlen_scalar[...] = np.array([1,2,3],np.int16)
        f.close()

    def tearDown(self):
        # Remove the temporary files
        os.remove(self.file)

    def runTest(self):
        """testing vlen variables"""
        f = Dataset(self.file, 'r')
        v = f.variables[VAR1_NAME]
        vs = f.variables[VAR2_NAME]
        vs_alt = f.variables[VAR3_NAME]
        assert list(f.vltypes.keys()) == [VL_NAME]
        assert f.vltypes[VL_NAME].dtype == VL_BASETYPE
        assert f.variables['string_scalar'][...] == 'foo'
        assert_array_equal(f.variables['vlen_scalar'][...],np.array([1,2,3],np.int16))
        data2 = v[:]
        data2s = vs[:]
        for i in range(nlons):
            for j in range(nlats):
                assert_array_equal(data2[j,i], data[j,i])
                assert datas[j,i] == data2s[j,i]
        assert_array_equal(datas, vs_alt[:])
        f.close()


class TestInvalidDataType(unittest.TestCase):
    def runTest(self):
        f = Dataset(FILE_NAME, 'w', format='NETCDF3_CLASSIC')
        f.createDimension('x', 1)
        # using assertRaisesRegext as a context manager
        # only works with python >= 2.7 (issue #497)
        #with self.assertRaisesRegexp(ValueError, 'strings are only supported'):
        #    f.createVariable('foo', str, ('x',))
        try:
            f.createVariable('foo', str, ('x',))
        except ValueError:
            pass
        f.close()
        os.remove(FILE_NAME)

class TestScalarVlenString(unittest.TestCase):
    # issue 333
    def runTest(self):
        f = Dataset(FILE_NAME, 'w', format='NETCDF4')
        teststring = f.createVariable('teststring', str)
        stringout = "yyyymmdd_hhmmss"
        teststring[()] = stringout
        f.close()
        f = Dataset(FILE_NAME)
        assert f.variables['teststring'][:] == stringout
        f.close()
        os.remove(FILE_NAME)

class TestIntegerIndex(unittest.TestCase):
    # issue 526
    def runTest(self):
        strtest = Dataset(FILE_NAME, 'w', format='NETCDF4')
        strtest.createDimension('tenstrings', 10)
        strtest.createVariable('tenstrings', str, ['tenstrings'])
        strtest['tenstrings'][np.int32(5)] = 'asdf'
        strtest['tenstrings'][6.0] = 'asdf'
        strtest.close()
        f = Dataset(FILE_NAME)
        assert f.variables['tenstrings'][np.int32(5)] == 'asdf'
        assert f.variables['tenstrings'][6.0] == 'asdf'
        f.close()
        os.remove(FILE_NAME)

class TestObjectArrayIndexing(unittest.TestCase):

    def setUp(self):
        self.file = FILE_NAME
        f = Dataset(self.file,'w')
        vlen_t = f.createVLType(VL_BASETYPE, VL_NAME)
        f.createDimension(DIM1_NAME,nlons)
        f.createDimension(DIM2_NAME,nlats)
        strings_alt = f.createVariable(VAR3_NAME, datas.astype(str).dtype,
                                       (DIM2_NAME, DIM1_NAME))
        strings_alt[:] = datas.astype(str)
        f.close()

    def tearDown(self):
        # Remove the temporary files
        os.remove(self.file)

    def runTest(self):
        """testing vlen variables"""
        f = Dataset(self.file, 'r')
        vs_alt = f.variables[VAR3_NAME]
        unicode_strings = vs_alt[:]
        fancy_indexed = unicode_strings[0][[1,2,4]]
        assert fancy_indexed[0] == 'abc'
        assert fancy_indexed[1] == 'abcd'
        assert fancy_indexed[2] == 'abcdef'
        f.close()

class VlenAppendTestCase(unittest.TestCase):
    def setUp(self):

        import netCDF4
        if netCDF4.__netcdf4libversion__ < "4.4.1":
            self.skip = True
            try:
                self.skipTest("This test requires NetCDF 4.4.1 or later.")
            except AttributeError:
                # workaround for Python 2.6 (skipTest(reason) is new
                # in Python 2.7)
                pass
        else:
            self.skip = False

        self.file = FILE_NAME
        f = Dataset(self.file, 'w')
        vlen_type = f.createVLType(np.float64, 'vltest')
        f.createDimension('x', None)
        v = f.createVariable('vl', vlen_type, 'x')
        w = f.createVariable('vl2', np.float64, 'x')
        f.close()

    def tearDown(self):
        # Remove the temporary files
        os.remove(self.file)

    def runTest(self):
        """testing appending to vlen variables (issue #527)."""
        # workaround for Python 2.6
        if self.skip:
            return

        f = Dataset(self.file, 'a')
        w = f.variables["vl2"]
        v = f.variables["vl"]
        w[0:3] = np.arange(3, dtype=np.float64)
        v[0]                    # sometimes crashes
        v[0].tolist()           # sometimes crashes
        v[0].size               # BOOM!
        f.close()

class Vlen_ScaledInts(unittest.TestCase):
    def setUp(self):
        self.file = FILE_NAME
        nc = Dataset(self.file, 'w')
        vlen_type = nc.createVLType(np.uint8, 'vltest')
        nc.createDimension('x', None)
        v = nc.createVariable('vl', vlen_type, 'x')
        v.scale_factor = 1./254.
        v.missing_value=np.array(255,np.uint8)
        # random lengths between 1 and 1000
        ilen = np.random.randint(1,1000,size=100)
        n = 0
        for nlen in ilen:
            data = np.random.uniform(low=0.0, high=1.0, size=nlen)
            v[n] = data
            if n==99: self.data = data
            n += 1
        nc.close()
    def tearDown(self):
        # Remove the temporary files
        os.remove(self.file)
    def runTest(self):
        """testing packing float vlens as scaled integers (issue #1003)."""
        nc = Dataset(self.file)
        data = nc['vl'][-1]
        # check max error of compression
        err = np.abs(data - self.data)
        assert(err.max() < nc['vl'].scale_factor)
        # turn off auto-scaling
        nc.set_auto_maskandscale(False)
        data = nc['vl'][-1]
        assert(data[-1] == np.around(self.data[-1]/nc['vl'].scale_factor))
        nc.close()

if __name__ == '__main__':
    unittest.main()

[((20, 7, 20, 35), 'numpy.empty', 'np.empty', ({(20, 16, 20, 27): 'nlats * nlons', (20, 28, 20, 34): 'object'}, {}), '(nlats * nlons, object)', True, 'import numpy as np\n'), ((21, 8, 21, 36), 'numpy.empty', 'np.empty', ({(21, 17, 21, 28): 'nlats * nlons', (21, 29, 21, 35): 'object'}, {}), '(nlats * nlons, object)', True, 'import numpy as np\n'), ((27, 7, 27, 37), 'numpy.reshape', 'np.reshape', ({(27, 18, 27, 22): 'data', (27, 23, 27, 36): '(nlats, nlons)'}, {}), '(data, (nlats, nlons))', True, 'import numpy as np\n'), ((28, 8, 28, 39), 'numpy.reshape', 'np.reshape', ({(28, 19, 28, 24): 'datas', (28, 25, 28, 38): '(nlats, nlons)'}, {}), '(datas, (nlats, nlons))', True, 'import numpy as np\n'), ((9, 12, 9, 67), 'tempfile.NamedTemporaryFile', 'tempfile.NamedTemporaryFile', (), '', False, 'import tempfile\n'), ((25, 14, 25, 45), 'numpy.arange', 'np.arange', (), '', True, 'import numpy as np\n'), ((228, 4, 228, 19), 'unittest.main', 'unittest.main', ({}, {}), '()', False, 'import unittest\n'), ((35, 12, 35, 34), 'netCDF4.Dataset', 'Dataset', ({(35, 20, 35, 29): 'self.file', (35, 30, 35, 33): '"""w"""'}, {}), "(self.file, 'w')", False, 'from netCDF4 import Dataset\n'), ((53, 27, 53, 53), 'numpy.array', 'np.array', ({(53, 36, 53, 43): '[1, 2, 3]', (53, 44, 53, 52): 'np.int16'}, {}), '([1, 2, 3], np.int16)', True, 'import numpy as np\n'), ((58, 8, 58, 28), 'os.remove', 'os.remove', ({(58, 18, 58, 27): 'self.file'}, {}), '(self.file)', False, 'import os\n'), ((62, 12, 62, 35), 'netCDF4.Dataset', 'Dataset', ({(62, 20, 62, 29): 'self.file', (62, 31, 62, 34): '"""r"""'}, {}), "(self.file, 'r')", False, 'from netCDF4 import Dataset\n'), ((76, 8, 76, 44), 'numpy.testing.assert_array_equal', 'assert_array_equal', ({(76, 27, 76, 32): 'datas', (76, 34, 76, 43): 'vs_alt[:]'}, {}), '(datas, vs_alt[:])', False, 'from numpy.testing import assert_array_equal\n'), ((82, 12, 82, 61), 'netCDF4.Dataset', 'Dataset', (), '', False, 'from netCDF4 import Dataset\n'), ((93, 8, 93, 28), 'os.remove', 'os.remove', ({(93, 18, 93, 27): 'FILE_NAME'}, {}), '(FILE_NAME)', False, 'import os\n'), ((98, 12, 98, 53), 'netCDF4.Dataset', 'Dataset', (), '', False, 'from netCDF4 import Dataset\n'), ((103, 12, 103, 30), 'netCDF4.Dataset', 'Dataset', ({(103, 20, 103, 29): 'FILE_NAME'}, {}), '(FILE_NAME)', False, 'from netCDF4 import Dataset\n'), ((106, 8, 106, 28), 'os.remove', 'os.remove', ({(106, 18, 106, 27): 'FILE_NAME'}, {}), '(FILE_NAME)', False, 'import os\n'), ((111, 18, 111, 59), 'netCDF4.Dataset', 'Dataset', (), '', False, 'from netCDF4 import Dataset\n'), ((117, 12, 117, 30), 'netCDF4.Dataset', 'Dataset', ({(117, 20, 117, 29): 'FILE_NAME'}, {}), '(FILE_NAME)', False, 'from netCDF4 import Dataset\n'), ((121, 8, 121, 28), 'os.remove', 'os.remove', ({(121, 18, 121, 27): 'FILE_NAME'}, {}), '(FILE_NAME)', False, 'import os\n'), ((127, 12, 127, 34), 'netCDF4.Dataset', 'Dataset', ({(127, 20, 127, 29): 'self.file', (127, 30, 127, 33): '"""w"""'}, {}), "(self.file, 'w')", False, 'from netCDF4 import Dataset\n'), ((138, 8, 138, 28), 'os.remove', 'os.remove', ({(138, 18, 138, 27): 'self.file'}, {}), '(self.file)', False, 'import os\n'), ((142, 12, 142, 35), 'netCDF4.Dataset', 'Dataset', ({(142, 20, 142, 29): 'self.file', (142, 31, 142, 34): '"""r"""'}, {}), "(self.file, 'r')", False, 'from netCDF4 import Dataset\n'), ((167, 12, 167, 35), 'netCDF4.Dataset', 'Dataset', ({(167, 20, 167, 29): 'self.file', (167, 31, 167, 34): '"""w"""'}, {}), "(self.file, 'w')", False, 'from netCDF4 import Dataset\n'), ((176, 8, 176, 28), 'os.remove', 'os.remove', ({(176, 18, 176, 27): 'self.file'}, {}), '(self.file)', False, 'import os\n'), ((184, 12, 184, 35), 'netCDF4.Dataset', 'Dataset', ({(184, 20, 184, 29): 'self.file', (184, 31, 184, 34): '"""a"""'}, {}), "(self.file, 'a')", False, 'from netCDF4 import Dataset\n'), ((187, 17, 187, 47), 'numpy.arange', 'np.arange', (), '', True, 'import numpy as np\n'), ((196, 13, 196, 36), 'netCDF4.Dataset', 'Dataset', ({(196, 21, 196, 30): 'self.file', (196, 32, 196, 35): '"""w"""'}, {}), "(self.file, 'w')", False, 'from netCDF4 import Dataset\n'), ((201, 24, 201, 46), 'numpy.array', 'np.array', ({(201, 33, 201, 36): '255', (201, 37, 201, 45): 'np.uint8'}, {}), '(255, np.uint8)', True, 'import numpy as np\n'), ((203, 15, 203, 49), 'numpy.random.randint', 'np.random.randint', (), '', True, 'import numpy as np\n'), ((213, 8, 213, 28), 'os.remove', 'os.remove', ({(213, 18, 213, 27): 'self.file'}, {}), '(self.file)', False, 'import os\n'), ((216, 13, 216, 31), 'netCDF4.Dataset', 'Dataset', ({(216, 21, 216, 30): 'self.file'}, {}), '(self.file)', False, 'from netCDF4 import Dataset\n'), ((219, 14, 219, 38), 'numpy.abs', 'np.abs', ({(219, 21, 219, 37): 'data - self.data'}, {}), '(data - self.data)', True, 'import numpy as np\n'), ((69, 59, 69, 85), 'numpy.array', 'np.array', ({(69, 68, 69, 75): '[1, 2, 3]', (69, 76, 69, 84): 'np.int16'}, {}), '([1, 2, 3], np.int16)', True, 'import numpy as np\n'), ((206, 19, 206, 66), 'numpy.random.uniform', 'np.random.uniform', (), '', True, 'import numpy as np\n'), ((224, 27, 224, 73), 'numpy.around', 'np.around', ({(224, 37, 224, 72): "(self.data[-1] / nc['vl'].scale_factor)"}, {}), "(self.data[-1] / nc['vl'].scale_factor)", True, 'import numpy as np\n'), ((74, 16, 74, 57), 'numpy.testing.assert_array_equal', 'assert_array_equal', ({(74, 35, 74, 45): 'data2[j, i]', (74, 47, 74, 56): 'data[j, i]'}, {}), '(data2[j, i], data[j, i])', False, 'from numpy.testing import assert_array_equal\n'), ((114, 30, 114, 41), 'numpy.int32', 'np.int32', ({(114, 39, 114, 40): '(5)'}, {}), '(5)', True, 'import numpy as np\n'), ((118, 41, 118, 52), 'numpy.int32', 'np.int32', ({(118, 50, 118, 51): '(5)'}, {}), '(5)', True, 'import numpy as np\n')]