add context manager support for IO (#3846)

* add context manager support for IO

add context manager support for IO

* update to use single quotes for consistency

Author: qindazhu

src/pybind/kaldi/table.py

@@ -91,6 +91,10 @@ def __init__(self, rspecifier=''):
     def __enter__(self):
         return self
 
+    def __exit__(self, type, value, traceback):
+        if self.IsOpen():
+            self.Close()
+
     def __iter__(self):
         while not self.Done():
             key = self.Key()
@@ -262,6 +266,10 @@ def __init__(self, rspecifier=''):
     def __enter__(self):
         return self
 
+    def __exit__(self, type, value, traceback):
+        if self.IsOpen():
+            self.Close()
+
     def __contains__(self, key):
         return self.HasKey(key)
 
@@ -417,6 +425,10 @@ def __init__(self, wspecifier=''):
     def __enter__(self):
         return self
 
+    def __exit__(self, type, value, traceback):
+        if self.IsOpen():
+            self.Close()
+    
     def __setitem__(self, key, value):
         self.Write(key, value)
 

src/pybind/tests/test_kaldi_pybind.py

@@ -58,6 +58,18 @@ def test_matrix_reader_writer(self):
         np.testing.assert_array_equal(value.numpy(), gold)
 
         matrix_reader.Close()
+
+        # test with context manager
+        kp_matrix[0, 0] = 20
+        with kaldi.MatrixWriter(wspecifier) as writer:
+            writer.Write('id_2', kp_matrix)
+        with kaldi.SequentialMatrixReader(rspecifier) as reader:
+            key = reader.Key()
+            self.assertEqual(key, 'id_2')
+            value = reader.Value()
+            gold = np.array([[20, 0, 0], [0, 0, 0]])
+            np.testing.assert_array_equal(value.numpy(), gold)
+        
         os.remove('test.ark')
 
     def test_matrix_reader_iterator(self):

src/pybind/tests/test_table_types.py

@@ -124,6 +124,13 @@ def test_float_matrix(self):
             np.testing.assert_array_almost_equal(
                 reader.Value(key).numpy(), data[key])
         reader.Close()
+        
+        # test RandomAccessReader with context manager        
+        with kaldi.RandomAccessMatrixReader(rspecifier) as reader:
+            for key in data.keys():
+                self.assertTrue(reader.HasKey(key))
+                np.testing.assert_array_almost_equal(
+                    reader.Value(key).numpy(), data[key])
 
         shutil.rmtree(tmp)