[Scipy-svn] r6867 - in trunk/scipy/io/matlab: . tests

[scipy-svn at scipy.org]

Author: matthew.brett at gmail.com
Date: 2010-11-11 18:56:45 -0600 (Thu, 11 Nov 2010)
New Revision: 6867

Modified:
   trunk/scipy/io/matlab/mio.py
   trunk/scipy/io/matlab/tests/test_mio.py
Log:
NF - added variable_names list to top level load API

Modified: trunk/scipy/io/matlab/mio.py
===================================================================
--- trunk/scipy/io/matlab/mio.py	2010-11-12 00:56:41 UTC (rev 6866)
+++ trunk/scipy/io/matlab/mio.py	2010-11-12 00:56:45 UTC (rev 6867)
@@ -123,6 +123,12 @@
     %(append_arg)s
     %(load_args)s
     %(struct_arg)s
+    variable_names : None or sequence
+        If None (the default) - read all variables in file. Otherwise
+        `variable_names` should be a sequence of strings, giving names of the
+        matlab variables to read from the file.  The reader will skip any
+        variable with a name not in this sequence, possibly saving some read
+        processing.
 
     Returns
     -------
@@ -138,8 +144,13 @@
     files.  Because scipy does not supply one, we do not implement the
     HDF5 / 7.3 interface here.
     '''
+    if 'variable_names' in kwargs:
+        variable_names = kwargs['variable_names']
+        del kwargs['variable_names']
+    else:
+        variable_names = None
     MR = mat_reader_factory(file_name, appendmat, **kwargs)
-    matfile_dict = MR.get_variables()
+    matfile_dict = MR.get_variables(variable_names)
     if mdict is not None:
         mdict.update(matfile_dict)
     else:

Modified: trunk/scipy/io/matlab/tests/test_mio.py
===================================================================
--- trunk/scipy/io/matlab/tests/test_mio.py	2010-11-12 00:56:41 UTC (rev 6866)
+++ trunk/scipy/io/matlab/tests/test_mio.py	2010-11-12 00:56:45 UTC (rev 6867)
@@ -800,5 +800,19 @@
     assert_equal(set(field_names), set(('a', 'b')))
 
 
+def test_loadmat_varnames():
+    # Test that we can get just one variable from a mat file using loadmat
+    eg_file = pjoin(test_data_path, 'testmulti_7.4_GLNX86.mat')
+    sys_v_names = ['__globals__',
+                   '__header__',
+                   '__version__']
+    vars = loadmat(eg_file)
+    assert_equal(set(vars.keys()), set(['a', 'theta'] + sys_v_names))
+    vars = loadmat(eg_file, variable_names=['a'])
+    assert_equal(set(vars.keys()), set(['a'] + sys_v_names))
+    vars = loadmat(eg_file, variable_names=['theta'])
+    assert_equal(set(vars.keys()), set(['theta'] + sys_v_names))
+
+
 if __name__ == "__main__":
     run_module_suite()