From 2da240c7ec776b44ffe6e06fa551fbff960c3b18 Mon Sep 17 00:00:00 2001
From: Helin Wang <helinwang@baidu.com>
Date: Mon, 14 Aug 2017 15:13:23 -0700
Subject: [PATCH] fix local recordio reader

---
 python/paddle/v2/reader/creator.py            |  12 ++++++----
 python/paddle/v2/reader/tests/creator_test.py |  22 ++++++++++++++++++
 .../v2/reader/tests/test_reader_recordio.dat  | Bin 0 -> 76 bytes
 3 files changed, 30 insertions(+), 4 deletions(-)
 create mode 100644 python/paddle/v2/reader/tests/test_reader_recordio.dat

diff --git a/python/paddle/v2/reader/creator.py b/python/paddle/v2/reader/creator.py
index d0f18e4b66..97e844b92c 100644
--- a/python/paddle/v2/reader/creator.py
+++ b/python/paddle/v2/reader/creator.py
@@ -57,7 +57,7 @@ def text_file(path):
     return reader
 
 
-def recordio_local(paths, buf_size=100):
+def recordio(paths, buf_size=100):
     """
     Creates a data reader from given RecordIO file paths separated by ",",
         glob pattern is supported.
@@ -67,15 +67,19 @@ def recordio_local(paths, buf_size=100):
 
     import recordio as rec
     import paddle.v2.reader.decorator as dec
+    import cPickle as pickle
 
     def reader():
-        a = ','.join(paths)
-        f = rec.reader(a)
+        if isinstance(paths, basestring):
+            path = paths
+        else:
+            path = ",".join(paths)
+        f = rec.reader(path)
         while True:
             r = f.read()
             if r is None:
                 break
-            yield r
+            yield pickle.loads(r)
         f.close()
 
     return dec.buffered(reader, buf_size)
diff --git a/python/paddle/v2/reader/tests/creator_test.py b/python/paddle/v2/reader/tests/creator_test.py
index 359f3eeefb..cf190aa664 100644
--- a/python/paddle/v2/reader/tests/creator_test.py
+++ b/python/paddle/v2/reader/tests/creator_test.py
@@ -34,5 +34,27 @@ class TestTextFile(unittest.TestCase):
             self.assertEqual(e, str(idx * 2) + " " + str(idx * 2 + 1))
 
 
+class TestRecordIO(unittest.TestCase):
+    def do_test(self, path):
+        reader = paddle.v2.reader.creator.recordio(path)
+        idx = 0
+        for e in reader():
+            if idx == 0:
+                self.assertEqual(e, (1, 2, 3))
+            elif idx == 1:
+                self.assertEqual(e, (4, 5, 6))
+            idx += 1
+        self.assertEqual(idx, 2)
+
+    def test_recordIO(self):
+        self.do_test(
+            os.path.join(
+                os.path.dirname(__file__), "test_reader_recordio.dat"))
+        self.do_test([
+            os.path.join(
+                os.path.dirname(__file__), "test_reader_recordio.dat")
+        ])
+
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/python/paddle/v2/reader/tests/test_reader_recordio.dat b/python/paddle/v2/reader/tests/test_reader_recordio.dat
new file mode 100644
index 0000000000000000000000000000000000000000..a99a35bb829e066c4845d0b85b96cd1eb3a12491
GIT binary patch
literal 76
zcmZQ!W@4P2Bs!asfq}sSh?#)+KN|x>v0q|9K_sIV14Bftj}1RiRKwGd%hQO<)0nHI
Tz>rH1B4onlY0Bkk1`z@P(}N7c

literal 0
HcmV?d00001

-- 
GitLab