|
1 | 1 | # benchmark reads and writes, with and without compression. |
2 | 2 | # tests all four supported file formats. |
| 3 | +from typing import TYPE_CHECKING, Any, Literal |
3 | 4 | from numpy.random.mtrand import uniform |
4 | 5 | import netCDF4 |
5 | 6 | from timeit import Timer |
6 | 7 | import os, sys |
| 8 | +if TYPE_CHECKING: |
| 9 | + from netCDF4 import Format as NCFormat |
| 10 | +else: |
| 11 | + NCFormat = Any |
7 | 12 |
|
8 | 13 | # create an n1dim by n2dim by n3dim random array. |
9 | 14 | n1dim = 30 |
|
14 | 19 | sys.stdout.write('reading and writing a %s by %s by %s by %s random array ..\n'%(n1dim,n2dim,n3dim,n4dim)) |
15 | 20 | array = uniform(size=(n1dim,n2dim,n3dim,n4dim)) |
16 | 21 |
|
17 | | -def write_netcdf(filename,zlib=False,least_significant_digit=None,format='NETCDF4',closeit=False): |
| 22 | +def write_netcdf(filename, zlib=False, least_significant_digit=None, format: NCFormat='NETCDF4',closeit=False): |
18 | 23 | file = netCDF4.Dataset(filename,'w',format=format,diskless=True,persist=True) |
19 | 24 | file.createDimension('n1', n1dim) |
20 | 25 | file.createDimension('n2', n2dim) |
@@ -42,13 +47,13 @@ def read_netcdf(ncfile): |
42 | 47 | sys.stdout.write('writing took %s seconds\n' %\ |
43 | 48 | repr(sum(t.repeat(ntrials,1))/ntrials)) |
44 | 49 | # test reading. |
45 | | - ncfile = write_netcdf('test1.nc',format=format) |
| 50 | + ncfile = write_netcdf('test1.nc',format=format) # type: ignore |
46 | 51 | t = Timer("read_netcdf(ncfile)","from __main__ import read_netcdf,ncfile") |
47 | 52 | sys.stdout.write('reading took %s seconds\n' % |
48 | 53 | repr(sum(t.repeat(ntrials,1))/ntrials)) |
49 | 54 |
|
50 | 55 | # test diskless=True in nc_open |
51 | | -format='NETCDF3_CLASSIC' |
| 56 | +format: Literal["NETCDF3_CLASSIC"] = 'NETCDF3_CLASSIC' # mypy should know this but it needs help... |
52 | 57 | trials=50 |
53 | 58 | sys.stdout.write('test caching of file in memory on open for %s\n' % format) |
54 | 59 | sys.stdout.write('testing file format %s ...\n' % format) |
|
0 commit comments