Added PyNIO tests. Added support to combine any WRF variable via lists and multiple times

wrf_open/var/src/python/wrf/var/__init__.py

@@ -159,7 +159,8 @@ _VALID_ARGS = {"cape2d" : ["missing", "timeidx"],
              "wspddir" : ["units", "timeidx"],
              "wspddir_uvmet" : ["units", "timeidx"],
              "wspddir_uvmet10" : ["units", "timeidx"],  
-             "ctt" : ["timeidx"]
+             "ctt" : ["timeidx"],
+             "default" : ["timeidx"]
             }
 
 _ALIASES = {"cape_2d" : "cape2d",
@@ -196,13 +197,20 @@ def _undo_alias(alias):
         return actual
 
 def _check_kargs(var, kargs):
-    for arg, val in kargs.iteritems():
+    for arg in kargs.iterkeys():
         if arg not in _VALID_ARGS[var]:
             raise ArgumentError("'%s' is an invalid keyword "
                           "argument for '%s" % (arg, var))
             
 
 def getvar(wrfnc, var, **kargs):
+    if is_standard_wrf_var(wrfnc, var):
+        if "timeidx" in kargs:
+            timeidx = kargs["timeidx"]
+        else:
+            timeidx = 0 
+        return extract_vars(wrfnc, timeidx, var)[var]
+    
     actual_var = _undo_alias(var)
     if actual_var not in _VALID_ARGS:
         raise ArgumentError("'%s' is not a valid variable name" % (var))

wrf_open/var/src/python/wrf/var/util.py

@@ -3,7 +3,7 @@ from collections import Iterable
 import numpy as n
 
 __all__ = ["extract_vars", "extract_global_attrs", "extract_dim",
-           "combine_files"]
+           "combine_files", "is_standard_wrf_var"]
 
 def _is_multi_time(timeidx):
     if timeidx == -1:
@@ -81,6 +81,12 @@ def extract_vars(wrfnc, timeidx, vars):
     else:
         return {var:combine_files(wrfnc, var, timeidx) for var in varlist}
     
+def is_standard_wrf_var(wrfnc, var):
+    multifile = _is_multi_file(wrfnc)
+    if multifile:
+        wrfnc = wrfnc[0]
+    return var in wrfnc.variables
+        
     
     
     

wrf_open/var/test/utests.py

@@ -5,13 +5,8 @@ import numpy.ma as ma
 import os, sys
 import subprocess
 
-import Ngl, Nio
 from wrf.var import getvar
 
-# This should work with Nio
-from netCDF4 import Dataset as NetCDF
-from boto.ec2.instancestatus import Status
-
 NCL_EXE = "/Users/ladwig/nclbuild/6.3.0/bin/ncl"
 TEST_FILE = "/Users/ladwig/Documents/wrf_files/wrfout_d01_2010-06-13_21:00:00"
 OUT_NC_FILE = "/tmp/wrftest.nc"
@@ -45,14 +40,40 @@ def setUpModule():
 
 # Using helpful information at: 
 # http://eli.thegreenplace.net/2014/04/02/dynamically-generating-python-test-cases
-def make_test(varname, wrf_in, referent, multi=False, repeat=3):
+def make_test(varname, wrf_in, referent, multi=False, repeat=3, pynio=False):
     def test(self):
+        try:
+            from netCDF4 import Dataset as NetCDF
+        except:
+            pass
+        
+        try:
+            from PyNIO import Nio
+        except:
+            pass
         
         if not multi:
+            if not pynio:
                 in_wrfnc = NetCDF(wrf_in)
             else:
+                # Note: Python doesn't like it if you reassign an outer scoped
+                # variable (wrf_in in this case)
+                if not wrf_in.endswith(".nc"):
+                    wrf_file = wrf_in + ".nc"
+                else:
+                    wrf_file = wrf_in
+                in_wrfnc = Nio.open_file(wrf_file)
+        else:
+            if not pynio:
                 nc = NetCDF(wrf_in)
                 in_wrfnc = [nc for i in xrange(repeat)]
+            else:
+                if not wrf_in.endswith(".nc"):
+                    wrf_file = wrf_in + ".nc"
+                else:
+                    wrf_file = wrf_in
+                nc = Nio.open_file(wrf_file)
+                in_wrfnc = [nc for i in xrange(repeat)]
         
         refnc = NetCDF(referent)
             
@@ -161,6 +182,11 @@ if __name__ == "__main__":
                 "theta", "tk", "tv", "twb", "updraft_helicity", "ua", "va", 
                 "wa", "uvmet10", "uvmet", "z", "ctt", "cape_2d", "cape_3d"]
     
+    try:
+        import netCDF4
+    except ImportError:
+        pass
+    else:
         for var in wrf_vars:
             if var in ignore_vars:
                 continue
@@ -170,6 +196,21 @@ if __name__ == "__main__":
             setattr(WRFVarsTest, 'test_{0}'.format(var), test_func1)
             setattr(WRFVarsTest, 'test_multi_{0}'.format(var), test_func2)
         
+    try:
+        import PyNIO
+    except ImportError:
+        pass
+    else:
+        for var in wrf_vars:
+            if var in ignore_vars:
+                continue
+            
+            test_func1 = make_test(var, TEST_FILE, OUT_NC_FILE,pynio=True)
+            test_func2 = make_test(var, TEST_FILE, OUT_NC_FILE, multi=True,
+                                   pynio=True)
+            setattr(WRFVarsTest, 'test_pynio_{0}'.format(var), test_func1)
+            setattr(WRFVarsTest, 'test_pynio_multi_{0}'.format(var), 
+                    test_func2)
         
     ut.main()