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

[scipy-svn at scipy.org]

Author: matthew.brett at gmail.com
Date: 2009-02-19 15:45:37 -0600 (Thu, 19 Feb 2009)
New Revision: 5576

Modified:
   trunk/scipy/io/matlab/mio5.py
   trunk/scipy/io/matlab/tests/test_mio.py
Log:
Added ability to pass dicts, objects for writing; test for 1D write/read shape

Modified: trunk/scipy/io/matlab/mio5.py
===================================================================
--- trunk/scipy/io/matlab/mio5.py	2009-02-19 20:42:17 UTC (rev 5575)
+++ trunk/scipy/io/matlab/mio5.py	2009-02-19 21:45:37 UTC (rev 5576)
@@ -879,6 +879,17 @@
 class Mat5WriterGetter(object):
     ''' Wraps stream and options, provides methods for getting Writer objects '''
     def __init__(self, stream, unicode_strings, long_field_names=False):
+        ''' Initialize writer getter
+
+        Parameters
+        ----------
+        stream : fileobj
+           object to which to write
+        unicode_strings : bool
+           If True, write unicode strings
+        long_field_names : bool, optional
+           If True, allow writing of long field names (127 bytes)
+        '''
         self.stream = stream
         self.unicode_strings = unicode_strings
         self.long_field_names = long_field_names
@@ -886,6 +897,68 @@
     def rewind(self):
         self.stream.seek(0)
 
+    def to_writeable(self, source):
+        ''' Convert input object ``source`` to something we can write
+
+        Parameters
+        ----------
+        source : object
+
+        Returns
+        -------
+        arr : ndarray
+
+        Examples
+        --------
+        >>> from StringIO import StringIO
+        >>> mwg = Mat5WriterGetter(StringIO(), True)
+        >>> mwg.to_writeable(np.array([1])) # pass through ndarrays
+        array([1])
+        >>> expected = np.array([(1, 2)], dtype=[('a', '|O8'), ('b', '|O8')])
+        >>> np.all(mwg.to_writeable({'a':1,'b':2}) == expected)
+        True
+        >>> class klass(object): pass
+        >>> c = klass
+        >>> c.a = 1
+        >>> c.b = 2
+        >>> np.all(mwg.to_writeable({'a':1,'b':2}) == expected)
+        True
+        >>> mwg.to_writeable([])
+        array([], dtype=float64)
+        >>> mwg.to_writeable(())
+        array([], dtype=float64)
+        >>> mwg.to_writeable(None)
+
+        >>> mwg.to_writeable('a string').dtype
+        dtype('|S8')
+        >>> mwg.to_writeable(1)
+        array(1)
+        >>> mwg.to_writeable([1])
+        array([1])
+        >>> mwg.to_writeable([1])
+        array([1])
+        >>> mwg.to_writeable(object()) # not convertable
+
+        '''
+        if isinstance(source, np.ndarray):
+            return source
+        if source is None:
+            return None
+        # Objects that have dicts
+        if hasattr(source, '__dict__'):
+            source = source.__dict__
+        # Mappings or object dicts
+        if hasattr(source, 'keys'):
+            dtype = [(k,object) for k in source]
+            return np.array( [tuple(source.itervalues())] ,dtype)
+        # Next try and convert to an array
+        narr = np.asanyarray(source)
+        if narr.dtype.type in (np.object, np.object_) and \
+           narr.shape == () and narr == source:
+            # No interesting conversion possible
+            return None
+        return narr
+
     def matrix_writer_factory(self, arr, name, is_global=False):
         ''' Factory function to return matrix writer given variable to write
 
@@ -901,11 +974,9 @@
         # First check if these are sparse
         if scipy.sparse.issparse(arr):
             return Mat5SparseWriter(self.stream, arr, name, is_global)
-        # Next try and convert to an array
-        narr = np.asanyarray(arr)
-        if narr.dtype.type in (np.object, np.object_) and \
-           narr.shape == () and narr == arr:
-            # No interesting conversion possible
+        # Try to convert things that aren't arrays
+        narr = self.to_writeable(arr)
+        if narr is None:
             raise TypeError('Could not convert %s (type %s) to array'
                             % (arr, type(arr)))
         args = (self.stream,

Modified: trunk/scipy/io/matlab/tests/test_mio.py
===================================================================
--- trunk/scipy/io/matlab/tests/test_mio.py	2009-02-19 20:42:17 UTC (rev 5575)
+++ trunk/scipy/io/matlab/tests/test_mio.py	2009-02-19 21:45:37 UTC (rev 5576)
@@ -13,6 +13,7 @@
 import gzip
 
 from numpy.testing import \
+     assert_array_equal, \
      assert_array_almost_equal, \
      assert_equal, \
      assert_raises
@@ -196,6 +197,8 @@
     'expected': {'testscalarobject': mlarr(1, dtype=object)}
     })
 '''
+
+
 def types_compatible(var1, var2):
     ''' Check if types are same or compatible
     
@@ -455,3 +458,22 @@
     sio.truncate(0)
     writer = Mat5NumericWriter(sio, np.zeros(10), 'aaaaaa').write()
     yield assert_true, sio.len - w_sz < 4
+
+
+def test_save_dict():
+    # Test that dict can be saved (as recarray), loaded as matstruct
+    d = {'a':1, 'b':2}
+    stream = StringIO()
+    savemat(stream, {'dict':d})
+    stream.seek(0)
+    vals = loadmat(stream)
+
+
+def test_1d_shape():
+    # Current behavior is 1D -> column vector
+    arr = np.arange(5)
+    stream = StringIO()
+    savemat(stream, {'oned':arr})
+    stream.seek(0)
+    vals = loadmat(stream)
+    assert_equal(vals['oned'].shape, (5,1))