diff --git a/python/paddle/v2/reader/creator.py b/python/paddle/v2/reader/creator.py index 07142056f872db5113acdd296b17c52b343c1be6..994062096fbe910c8189d3ecd49298781dec708f 100644 --- a/python/paddle/v2/reader/creator.py +++ b/python/paddle/v2/reader/creator.py @@ -16,7 +16,7 @@ Creator package contains some simple reader creator, which could be used in user program. """ -__all__ = ['np_array', 'text_file'] +__all__ = ['np_array', 'text_file', "RecordIO"] def np_array(x): @@ -55,3 +55,22 @@ def text_file(path): f.close() return reader + + +def RecordIO(path): + """ + Creates a data reader that outputs record one one by one from given recordio file + :path: path of recordio file + :returns: data reader of recordio file + """ + + def reader(): + f = recordio.reader(path) + while True: + r = f.read() + if r is None: + break + yield r + f.close() + + return reader diff --git a/python/paddle/v2/reader/tests/creator_test.py b/python/paddle/v2/reader/tests/creator_test.py index 9f8d7133b8694aae5541eff9576eaba8a31e77dc..dd84fbb002543f2c3d90c9543e475516165a5f4d 100644 --- a/python/paddle/v2/reader/tests/creator_test.py +++ b/python/paddle/v2/reader/tests/creator_test.py @@ -36,5 +36,14 @@ class TestTextFile(unittest.TestCase): self.assertEqual(e, str(idx * 2) + " " + str(idx * 2 + 1)) +class TestRecordIO(unittest.TestCase): + def test_RecordIO(self): + path = os.path.join( + os.path.dirname(__file__), "test_recordio_creator.dat") + reader = paddle.v2.reader.creator.RecordIO(path) + for idx, r in enumerate(reader()): + self.assertSequenceEqual(r, str(idx)) + + if __name__ == '__main__': unittest.main() diff --git a/python/paddle/v2/reader/tests/test_recordio_creator.dat b/python/paddle/v2/reader/tests/test_recordio_creator.dat new file mode 100644 index 0000000000000000000000000000000000000000..17aa89b6796184407e83246d3f342a55a66b4a69 Binary files /dev/null and b/python/paddle/v2/reader/tests/test_recordio_creator.dat differ