From da96e4771797b1a823be8fd5f7dcb4287b1eecbb Mon Sep 17 00:00:00 2001 From: Suizer98 Date: Fri, 12 Jan 2024 12:47:56 +0800 Subject: [PATCH] Change functionality of read_array so it returns multiple variables --- nco/nco.py | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/nco/nco.py b/nco/nco.py index 31574bc..1aa501b 100644 --- a/nco/nco.py +++ b/nco/nco.py @@ -495,15 +495,20 @@ def open_cdf(self, infile): return file_obj - def read_array(self, infile, var_name): - """Directly return a numpy array for a given variable name""" + def read_array(self, infile, var_names): + """Directly return multiple numpy arrays for given variable names""" file_handle = self.read_cdf(infile) - try: - # return the data array - return file_handle.variables[var_name][:] - except KeyError: - print("Cannot find variable: {0}".format(var_name)) - raise KeyError + result = {} + + for var_name in var_names: + try: + # return the data array for each variable + result[var_name] = file_handle.variables[var_name][:] + except KeyError: + print("Cannot find variable: {0}".format(var_name)) + raise KeyError + + return result def read_ma_array(self, infile, var_name): """Create a masked array based on cdf's FillValue"""