From 57b64d0d1902eb51bf79f595626c2b9f80a9d1e2 Mon Sep 17 00:00:00 2001
From: Josh Rosen <joshrosen@eecs.berkeley.edu>
Date: Fri, 1 Feb 2013 00:25:19 -0800
Subject: Fix stdout redirection in PySpark.

---
 python/pyspark/tests.py  | 9 +++++++++
 python/pyspark/worker.py | 5 +++--
 2 files changed, 12 insertions(+), 2 deletions(-)

(limited to 'python')

diff --git a/python/pyspark/tests.py b/python/pyspark/tests.py
index df7235756d..52297d44e6 100644
--- a/python/pyspark/tests.py
+++ b/python/pyspark/tests.py
@@ -108,5 +108,14 @@ class TestAddFile(PySparkTestCase):
         self.assertEqual("Hello World!", UserClass().hello())
 
 
+class TestIO(PySparkTestCase):
+
+    def test_stdout_redirection(self):
+        import subprocess
+        def func(x):
+            subprocess.check_call('ls', shell=True)
+        self.sc.parallelize([1]).foreach(func)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 9622e0cfe4..812e7a9da5 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -1,6 +1,7 @@
 """
 Worker that receives input from Piped RDD.
 """
+import os
 import sys
 import traceback
 from base64 import standard_b64decode
@@ -15,8 +16,8 @@ from pyspark.serializers import write_with_length, read_with_length, write_int,
 
 
 # Redirect stdout to stderr so that users must return values from functions.
-old_stdout = sys.stdout
-sys.stdout = sys.stderr
+old_stdout = os.fdopen(os.dup(1), 'w')
+os.dup2(2, 1)
 
 
 def load_obj():
-- 
cgit v1.2.3