From dbfe58ca47e8ff59d8d020739529ed6f02f5b9f0 Mon Sep 17 00:00:00 2001
From: gongweibao <weibao.gong@gmail.com>
Date: Wed, 21 Jun 2017 19:27:12 +0800
Subject: [PATCH] add local recordio reader interface

---
 python/paddle/v2/reader/creator.py            |  21 +++++++++++++++++-
 python/paddle/v2/reader/tests/creator_test.py |   9 ++++++++
 .../v2/reader/tests/test_recordio_creator.dat | Bin 0 -> 88 bytes
 3 files changed, 29 insertions(+), 1 deletion(-)
 create mode 100644 python/paddle/v2/reader/tests/test_recordio_creator.dat

diff --git a/python/paddle/v2/reader/creator.py b/python/paddle/v2/reader/creator.py
index 07142056f..994062096 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 9f8d7133b..dd84fbb00 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
GIT binary patch
literal 88
zcmZQ!W@2QOHw<B9U|?_oVlE*5&&I$|?3Wl&5Xor9z;M0c)+Lav0f;aJ5k?@w7(|$W
R2vZPY1|rNsgawGO1OWMk36uZ;

literal 0
HcmV?d00001

-- 
GitLab