From 731da9d48881a05247539b210786075548ba8aa5 Mon Sep 17 00:00:00 2001
From: Bill Ladwig <ladwig@ucar.edu>
Date: Sat, 12 Dec 2015 13:23:14 -0700
Subject: [PATCH] Added PyNIO tests.  Added support to combine any WRF variable
 via lists and multiple times

---
 wrf_open/var/src/python/wrf/__init__.pyc    | Bin 0 -> 373 bytes
 wrf_open/var/src/python/wrf/var/__init__.py |  12 ++-
 wrf_open/var/src/python/wrf/var/util.py     |   8 +-
 wrf_open/var/test/utests.py                 |  77 +++++++++++++++-----
 4 files changed, 76 insertions(+), 21 deletions(-)
 create mode 100644 wrf_open/var/src/python/wrf/__init__.pyc

diff --git a/wrf_open/var/src/python/wrf/__init__.pyc b/wrf_open/var/src/python/wrf/__init__.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3228da818fb46bcbff3227ef656bb45f1abf2c78
GIT binary patch
literal 373
zcmYLDK~BRk5F9sYfj|#^0Qa2m2N2=_2QEEO5LG3#ijt*Glh~2%P<rN5{EG*`CW_RR
zJf0bk$BusI$<M|6OOA1-vfs<{D}m!R3ZPJ^fepaGUV$h0R#0z7BjyNI0=I-xYD~u1
z#NZhFDR?%7C)mtj6xLF-N+;4={>d!UGofKF&{kV=>}}_A_F*AT!FioICX+Sn+brjB
zE-nq4C|rtoYFg{UqjQ3NlazF!tcNrH3Tz5$v#=X4_LL}<98Tb$Me0}2&g(iW4&_!q
z*t~1l1g{U)Rm0Qz;6BK<Y)<4ebK1LHx5u!vgRhB7QwE|N@%q2sHn!+$zEi?|A`q!7
IwNP`xe=rGLQ~&?~

literal 0
HcmV?d00001

diff --git a/wrf_open/var/src/python/wrf/var/__init__.py b/wrf_open/var/src/python/wrf/var/__init__.py
index 12a5a0a..42d1417 100755
--- a/wrf_open/var/src/python/wrf/var/__init__.py
+++ b/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))
diff --git a/wrf_open/var/src/python/wrf/var/util.py b/wrf_open/var/src/python/wrf/var/util.py
index 892d943..568b472 100644
--- a/wrf_open/var/src/python/wrf/var/util.py
+++ b/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:
@@ -80,6 +80,12 @@ def extract_vars(wrfnc, timeidx, vars):
         return {var:wrfnc.variables[var][:] for var in varlist}
     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
         
     
     
diff --git a/wrf_open/var/test/utests.py b/wrf_open/var/test/utests.py
index 776f025..cc7cfbc 100644
--- a/wrf_open/var/test/utests.py
+++ b/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:
-            in_wrfnc = NetCDF(wrf_in)
+            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:
-            nc = NetCDF(wrf_in)
-            in_wrfnc = [nc for i in xrange(repeat)]
+            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,15 +182,35 @@ if __name__ == "__main__":
                 "theta", "tk", "tv", "twb", "updraft_helicity", "ua", "va", 
                 "wa", "uvmet10", "uvmet", "z", "ctt", "cape_2d", "cape_3d"]
     
-    for var in wrf_vars:
-        if var in ignore_vars:
-            continue
+    try:
+        import netCDF4
+    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)
+            test_func2 = make_test(var, TEST_FILE, OUT_NC_FILE, multi=True)
+            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)
         
-        test_func1 = make_test(var, TEST_FILE, OUT_NC_FILE)
-        test_func2 = make_test(var, TEST_FILE, OUT_NC_FILE, multi=True)
-        setattr(WRFVarsTest, 'test_{0}'.format(var), test_func1)
-        setattr(WRFVarsTest, 'test_multi_{0}'.format(var), test_func2)
-    
-    
     ut.main()
     
\ No newline at end of file