From a4ff01db7a602c1c15b1383380ea59c663cf0eb8 Mon Sep 17 00:00:00 2001
From: luke-zhu <luke_zhu@brown.edu>
Date: Tue, 20 Feb 2018 01:52:03 -0500
Subject: [PATCH] Revert invalid use of io.StringIO in utils/profiler.py

---
 sdks/python/apache_beam/utils/profiler.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/sdks/python/apache_beam/utils/profiler.py b/sdks/python/apache_beam/utils/profiler.py
index 74d6abc600f5..9f9c8cd16296 100644
--- a/sdks/python/apache_beam/utils/profiler.py
+++ b/sdks/python/apache_beam/utils/profiler.py
@@ -21,15 +21,20 @@
 """
 
 import cProfile
-import io
 import logging
 import os
 import pstats
+import sys
 import tempfile
 import time
 import warnings
 from threading import Timer
 
+if sys.version_info[0] < 3:
+  import StringIO
+else:
+  from io import StringIO
+
 
 class Profile(object):
   """cProfile wrapper context for saving and logging profiler results."""
@@ -66,7 +71,7 @@ def __exit__(self, *args):
       os.remove(filename)
 
     if self.log_results:
-      s = io.StringIO()
+      s = StringIO()
       self.stats = pstats.Stats(
           self.profile, stream=s).sort_stats(Profile.SORTBY)
       self.stats.print_stats()