Merge pull request #15 from synthicity/updates-from-activitysim

New tests and bug fixes

api/python/openmatrix/File.py

@@ -27,10 +27,10 @@ def createMatrix(self, name, atom=None, shape=None, title='', filters=None,
            an existing numpy matrix, or a shape and an atom type."""
 
         # If object was passed in, make sure its shape is correct
-        if self.shape() != None and obj != None and obj.shape != self.shape():
+        if self.shape() is not None and obj is not None and obj.shape != self.shape():
             raise ShapeError('%s has shape %s but this file requires shape %s' %
                 (name, obj.shape, self.shape()))
-                
+
         # Create the HDF5 array
         if tables.__version__.startswith('3'):
             matrix = self.createCArray(self.root.data, name, atom, shape, title, filters,
@@ -39,11 +39,11 @@ def createMatrix(self, name, atom=None, shape=None, title='', filters=None,
             # this version is tables 2.4-compatible:
             matrix = self.createCArray(self.root.data, name, atom, shape, title, filters,
                                        chunkshape, byteorder, createparents)
-            if (obj != None):
+            if (obj is not None):
                 matrix[:] = obj
 
         # Store shape if we don't have one yet
-        if self._shape == None:
+        if self._shape is None:
             storeshape = np.array([matrix.shape[0],matrix.shape[1]], dtype='int32')
             self.root._v_attrs['SHAPE'] = storeshape
             self._shape = matrix.shape
@@ -96,10 +96,9 @@ def listMatrices(self):
     def listAllAttributes(self):
         """Return combined list of all attributes used for any Matrix in this File"""
         all_tags = set()
-        for m in self.listNodes(self.root,'CArray'):
-            if m.attrs != None:
-                all_tags.update(m.attrs._v_attrnames)
-        return sorted(list(all_tags))
+        for m in self.iterNodes(self.root.data, 'CArray'):
+            all_tags.update(m.attrs._v_attrnamesuser)
+        return sorted(all_tags)
 
 
     # MAPPINGS -----------------------------------------------
@@ -199,11 +198,12 @@ def _getMatricesByAttribute(self, key, value, matrices=None):
 
         answer = []
 
-        if matrices==None:
+        if matrices is None:
             matrices = self.listNodes(self.root.data,'CArray')
 
         for m in matrices:
-            if m.attrs == None: continue
+            if m.attrs is None:
+                continue
 
             # Only test if key is present in matrix attributes
             if key in m.attrs._v_attrnames and m.attrs[key] == value:

api/python/openmatrix/test/test_file.py

@@ -1,79 +1,184 @@
-import tables,numpy,os
+import os
+import tempfile
+
+import numpy as np
+import numpy.testing as npt
+import tables
+
 import openmatrix as omx
-from nose.tools import *
 
-def setup_clean():
-    try:
-        os.remove('test1.omx')
-    except:
+import nose.tools as nt
+
+TEST_FILE = None
+
+
+def ones5x5():
+    return np.ones((5, 5))
+
+
+def add_m1_node(f):
+    f.createMatrix('m1', obj=ones5x5())
+
+
+def setup_func():
+    global TEST_FILE
+
+    if TEST_FILE is not None and os.path.isfile(TEST_FILE):
+        os.remove(TEST_FILE)
+
+    with tempfile.NamedTemporaryFile(suffix='.omx') as tmp:
+        TEST_FILE = tmp.name
+
+
+def teardown_func():
+    if TEST_FILE is not None and os.path.isfile(TEST_FILE):
+        os.remove(TEST_FILE)
+
+
+@nt.with_setup(setup_func, teardown_func)
+def test_create_file():
+    with omx.openFile(TEST_FILE, 'w'):
         pass
+    assert os.path.isfile(TEST_FILE)
 
 
-def setup_empty_file():
-    try:
-        os.remove('test1.omx')
-        f = omx.openFile('test.omx','w')
-        f.close()
-    except:
+@nt.with_setup(setup_func, teardown_func)
+def test_open_readonly_hdf5_file():
+    with tables.openFile(TEST_FILE, 'w'):
         pass
 
+    assert os.path.isfile(TEST_FILE)
 
-def teardown_clean():
-    try:
-        os.remove('test2.omx')
-    except:
+    with omx.openFile(TEST_FILE, 'r'):
         pass
 
 
-@with_setup(setup_clean,teardown_clean)
-def test_create_file():
-    f = omx.openFile('test1.omx','w')
-    f.close()
-    assert(os.path.exists('test1.omx'))
+@nt.with_setup(setup_func, teardown_func)
+def test_set_get_del():
+    with omx.openFile(TEST_FILE, 'w') as f:
+        add_m1_node(f)
+        npt.assert_array_equal(f['m1'], ones5x5())
+        nt.assert_equal(f.shape(), (5, 5))
+        del f['m1']
+        nt.assert_not_in('m1', f)
 
-def test_open_readonly_hdf5_file():
-    f = tables.openFile('test2.omx','w')
-    f.close()
-    f = omx.openFile('test2.omx','r')
-    f.close() 
 
+@nt.with_setup(setup_func, teardown_func)
 def test_add_numpy_matrix_using_brackets():
-    f = omx.openFile('test3.omx','w')
-    f['m1'] = numpy.ones((5,5))
-    f.close()
+    with omx.openFile(TEST_FILE, 'w') as f:
+        f['m1'] = ones5x5()
+        npt.assert_array_equal(f['m1'], ones5x5())
+        nt.assert_equal(f.shape(), (5, 5))
+
+        # test check for shape matching
+        with nt.assert_raises(omx.Exceptions.ShapeError):
+            f.createMatrix('m2', obj=np.ones((8, 8)))
+
 
+@nt.with_setup(setup_func, teardown_func)
 def test_add_numpy_matrix_using_create_matrix():
-    f = omx.openFile('test4.omx','w')
-    f.createMatrix('m1', obj=numpy.ones((5,5)))
-    f.close()
+    with omx.openFile(TEST_FILE, 'w') as f:
+        f.createMatrix('m1', obj=ones5x5())
+        npt.assert_array_equal(f['m1'], ones5x5())
+        nt.assert_equal(f.shape(), (5, 5))
 
+
+@nt.with_setup(setup_func, teardown_func)
+@nt.raises(tables.FileModeError)
 def test_add_matrix_to_readonly_file():
-    f = omx.openFile('test6.omx','w')
-    f['m2'] = numpy.ones((5,5))
-    f.close()
-    f = omx.openFile('test6.omx','r')
-    assert_raises(tables.FileModeError, add_m1_node, f)
-    f.close() 
+    with omx.openFile(TEST_FILE, 'w') as f:
+        f['m2'] = np.ones((5, 5))
+
+    with omx.openFile(TEST_FILE, 'r') as f:
+        f.createMatrix('m1', obj=np.ones((5, 5)))
+
 
+@nt.with_setup(setup_func, teardown_func)
+@nt.raises(tables.NodeError)
 def test_add_matrix_with_same_name():
-    f = omx.openFile('test5.omx','w')
-    add_m1_node(f)
-    # now add m1 again:
-    assert_raises(tables.NodeError, add_m1_node, f)
-    f.close()
+    with omx.openFile(TEST_FILE, 'w') as f:
+        add_m1_node(f)
+        # now add m1 again:
+        add_m1_node(f)
 
-@with_setup(setup_clean,teardown_clean)
+
+@nt.with_setup(setup_func, teardown_func)
 def test_get_length_of_file():
-    f = omx.openFile('test7.omx','w')
-    f['m1'] = numpy.ones((5,5))
-    f['m2'] = numpy.ones((5,5))
-    f['m3'] = numpy.ones((5,5))
-    f['m4'] = numpy.ones((5,5))
-    f['m5'] = numpy.ones((5,5))
-    assert(len(f)==5)
-    assert(len(f.listMatrices())==5)
-    f.close()
+    with omx.openFile(TEST_FILE, 'w') as f:
+        f['m1'] = np.ones((5, 5))
+        f['m2'] = np.ones((5, 5))
+        f['m3'] = np.ones((5, 5))
+        f['m4'] = np.ones((5, 5))
+        f['m5'] = np.ones((5, 5))
+        nt.assert_equal(len(f), 5)
+        nt.assert_equal(len(f.listMatrices()), 5)
+
+
+@nt.with_setup(setup_func, teardown_func)
+def test_len_list_iter():
+    names = ['m{}'.format(x) for x in range(5)]
+    with omx.openFile(TEST_FILE, 'w') as f:
+        for m in names:
+            f[m] = ones5x5()
+
+        for mat in f:
+            npt.assert_array_equal(mat, ones5x5())
+
+        nt.assert_equal(len(f), len(names))
+        nt.assert_equal(f.listMatrices(), names)
+
+
+@nt.with_setup(setup_func, teardown_func)
+def test_contains():
+    with omx.openFile(TEST_FILE, 'w') as f:
+        add_m1_node(f)
+        nt.assert_in('m1', f)
+        # keep this here to be sure we're actually running
+        # File.__contains__
+        assert 'm1' in f
+
+
+@nt.with_setup(setup_func, teardown_func)
+def test_list_all_attrs():
+    with omx.openFile(TEST_FILE, 'w') as f:
+        add_m1_node(f)
+        f['m2'] = ones5x5()
+
+        nt.assert_equal(f.listAllAttributes(), [])
+
+        f['m1'].attrs['a1'] = 'a1'
+        f['m1'].attrs['a2'] = 'a2'
+        f['m2'].attrs['a2'] = 'a2'
+        f['m2'].attrs['a3'] = 'a3'
+
+        nt.assert_equal(f.listAllAttributes(), ['a1', 'a2', 'a3'])
+
+
+@nt.with_setup(setup_func, teardown_func)
+def test_matrices_by_attr():
+    with omx.openFile(TEST_FILE, 'w') as f:
+        f['m1'] = ones5x5()
+        f['m2'] = ones5x5()
+        f['m3'] = ones5x5()
+
+        for m in f:
+            m.attrs['a1'] = 'a1'
+            m.attrs['a2'] = 'a2'
+        f['m3'].attrs['a2'] = 'a22'
+        f['m3'].attrs['a3'] = 'a3'
+
+        gmba = f._getMatricesByAttribute
+
+        nt.assert_equal(gmba('zz', 'zz'), [])
+        nt.assert_equal(gmba('a1', 'a1'), [f['m1'], f['m2'], f['m3']])
+        nt.assert_equal(gmba('a2', 'a2'), [f['m1'], f['m2']])
+        nt.assert_equal(gmba('a2', 'a22'), [f['m3']])
+        nt.assert_equal(gmba('a3', 'a3'), [f['m3']])
 
-def add_m1_node(f):
-    f.createMatrix('m1', obj=numpy.ones((7,7)))
 
+@nt.with_setup(setup_func, teardown_func)
+def test_set_with_carray():
+    with omx.openFile(TEST_FILE, 'w') as f:
+        f['m1'] = ones5x5()
+        f['m2'] = f['m1']
+        npt.assert_array_equal(f['m2'], f['m1'])