from netCDF4 import Dataset, MFDataset, MFTime

import numpy as np

from numpy.random import seed, randint

from numpy.testing import assert_array_equal, assert_equal

from numpy import ma

import tempfile, unittest, os, datetime

import cftime

from pkg_resources import parse_version

nx=100; ydim=5; zdim=10

nfiles = 10

ninc = nx/nfiles

files = [tempfile.NamedTemporaryFile(suffix='.nc', delete=False).name for nfile in range(nfiles)]

data = randint(0,10,size=(nx,ydim,zdim))

missval = 99

data[::10] = missval

data = ma.masked_values(data,missval)

class VariablesTestCase(unittest.TestCase):

def setUp(self):

self.files = files

for nfile,file in enumerate(self.files):

f = Dataset(file,'w',format='NETCDF4_CLASSIC')

#f.createDimension('x',None)

f.createDimension('x',ninc)

f.createDimension('y',ydim)

f.createDimension('z',zdim)

f.history = 'created today'

x = f.createVariable('x','i',('x',))

x.units = 'zlotys'

dat = f.createVariable('data','i',('x','y','z',))

dat.long_name = 'phony data'

dat.missing_value = missval

nx1 = int(nfile*ninc); nx2 = int(ninc*(nfile+1))

#x[0:ninc] = np.arange(nfile*ninc,ninc*(nfile+1))

x[:] = np.arange(nfile*ninc,ninc*(nfile+1))

#dat[0:ninc] = data[nx1:nx2]

dat[:] = data[nx1:nx2]

f.close()

def tearDown(self):

# Remove the temporary files

for file in self.files:

os.remove(file)

def runTest(self):

"""testing multi-file dataset access"""

# specify the aggregation dim (not necessarily unlimited)

f = MFDataset(self.files,aggdim='x',check=True)

assert f.history == 'created today'

assert_array_equal(np.arange(0,nx),f.variables['x'][:])

varin = f.variables['data']

datin = varin[:]

assert_array_equal(datin.mask,data.mask)

varin.set_auto_maskandscale(False)

data2 = data.filled()

assert varin.long_name == 'phony data'

assert len(varin) == nx

assert varin.shape == (nx,ydim,zdim)

assert varin.dimensions == ('x','y','z')

assert_array_equal(varin[4:-4:4,3:5,2:8],data2[4:-4:4,3:5,2:8])

assert varin[0,0,0] == data2[0,0,0]

assert_array_equal(varin[:],data2)

assert getattr(varin,'nonexistantatt',None) == None

f.close()

class NonuniformTimeTestCase(unittest.TestCase):

ninc = 365

def setUp(self):

self.files = [tempfile.NamedTemporaryFile(suffix='.nc', delete=False).name for nfile in range(2)]

for nfile,file in enumerate(self.files):

f = Dataset(file,'w',format='NETCDF4_CLASSIC')

f.createDimension('time',None)

f.createDimension('y',ydim)

f.createDimension('z',zdim)

f.history = 'created today'

time = f.createVariable('time', 'f', ('time', ))

#time.units = 'days since {0}-01-01'.format(1979+nfile)

yr = 1979+nfile

time.units = 'days since %s-01-01' % yr

time.calendar = 'standard'

x = f.createVariable('x','f',('time', 'y', 'z'))

x.units = 'potatoes per square mile'

nx1 = self.ninc*nfile;

nx2 = self.ninc*(nfile+1)

time[:] = np.arange(self.ninc)

x[:] = np.arange(nx1, nx2).reshape(self.ninc,1,1) * np.ones((1, ydim, zdim))

f.close()

def tearDown(self):

# Remove the temporary files

for file in self.files:

os.remove(file)

def runTest(self):

# Get the real dates

# skip this until cftime pull request #55 is in a released

# version (1.0.1?). Otherwise, fix for issue #808 breaks this

if parse_version(cftime.__version__) >= parse_version('1.0.1'):

dates = []

for file in self.files:

f = Dataset(file)

t = f.variables['time']

dates.extend(cftime.num2date(t[:], t.units, t.calendar))

f.close()

# Compare with the MF dates

f = MFDataset(self.files,check=True)

t = f.variables['time']

mfdates = cftime.num2date(t[:], t.units, t.calendar)

T = MFTime(t)

assert_equal(len(T), len(t))

assert_equal(T.shape, t.shape)

assert_equal(T.dimensions, t.dimensions)

assert_equal(T.typecode(), t.typecode())

# skip this until cftime pull request #55 is in a released

# version (1.0.1?). Otherwise, fix for issue #808 breaks this

if parse_version(cftime.__version__) >= parse_version('1.0.1'):

assert_array_equal(cftime.num2date(T[:], T.units, T.calendar), dates)

assert_equal(cftime.date2index(datetime.datetime(1980, 1, 2), T), 366)

f.close()

if __name__ == '__main__':

unittest.main()