FIX: Minor fixes

larsoner · larsoner · commit f1a98fb4ff0d · 2016-03-14T11:35:47.000-04:00
diff --git a/sphinx_gallery/gen_rst.py b/sphinx_gallery/gen_rst.py
@@ -12,9 +12,11 @@
 Files that generate images should start with 'plot'
 
 """
-from __future__ import division, print_function, absolute_import
+from __future__ import (division, print_function, absolute_import,
+                        unicode_literals)
 from time import time
 import ast
+import codecs
 import hashlib
 import os
 import re
@@ -49,9 +51,9 @@ def prefixed_lines():
         return ''.join(prefixed_lines())
 
 try:
-    from StringIO import StringIO
+    from BytesIO import BytesIO
 except ImportError:
-    from io import StringIO
+    from io import BytesIO
 
 try:
     # make sure that the Agg backend is set before importing any
@@ -87,13 +89,23 @@ def __init__(self, file1, file2):
 
     def write(self, data):
         self.file1.write(data)
-        self.file2.write(unicode(data))  # this is the StringIO
+        self.file2.write(data)
 
     def flush(self):
         self.file1.flush()
         self.file2.flush()
 
 
+class MyBytesIO(BytesIO):
+    """Helper to deal with unicode not having buffer interface on Py2.7"""
+    def write(self, data):
+        # "unicode" objects don't have a buffer interface on Py2.7,
+        # so we need to manually convert
+        if isinstance(data, unicode):
+            data = data.encode('utf-8')
+        super(MyBytesIO, self).write(data)
+
+
 ###############################################################################
 CODE_DOWNLOAD = """**Total running time of the script:**
 ({0:.0f} minutes {1:.3f} seconds)\n\n
@@ -125,7 +137,7 @@ def flush(self):
 """
 
 
-CODE_OUTPUT = u""".. rst-class:: sphx-glr-script-out
+CODE_OUTPUT = """.. rst-class:: sphx-glr-script-out
 
  Out::
 
@@ -144,6 +156,9 @@ def get_docstring_and_rest(filename):
     rest: str
         `filename` content without the docstring
     """
+    # can't use codecs.open(filename, 'r', 'utf-8') here b/c ast doesn't
+    # seem to work with unicode strings in Python2.7
+    # "SyntaxError: encoding declaration in Unicode string"
     with open(filename, 'rb') as f:
         content = f.read()
 
@@ -206,8 +221,8 @@ def split_code_and_text_blocks(source_file):
 
 def codestr2rst(codestr, lang='python'):
     """Return reStructuredText code block from code string"""
-    code_directive = u"\n.. code-block:: {0}\n\n".format(lang)
-    indented_block = indent(codestr, u' ' * 4)
+    code_directive = "\n.. code-block:: {0}\n\n".format(lang)
+    indented_block = indent(codestr, ' ' * 4)
     return code_directive + indented_block
 
 
@@ -464,7 +479,7 @@ def execute_script(code_block, example_globals, image_path, fig_count,
         # First cd in the original example dir, so that any file
         # created by the example get created in this directory
         os.chdir(os.path.dirname(src_file))
-        my_buffer = StringIO()
+        my_buffer = MyBytesIO()
         my_stdout = Tee(sys.stdout, my_buffer)
         sys.stdout = my_stdout
 
@@ -474,9 +489,9 @@ def execute_script(code_block, example_globals, image_path, fig_count,
 
         sys.stdout = orig_stdout
 
-        my_stdout = my_buffer.getvalue().strip().expandtabs()
+        my_stdout = my_buffer.getvalue().decode('utf-8').strip().expandtabs()
         if my_stdout:
-            stdout = CODE_OUTPUT.format(indent(my_stdout, u' ' * 4))
+            stdout = CODE_OUTPUT.format(indent(my_stdout, ' ' * 4))
         os.chdir(cwd)
         figure_list = save_figures(image_path, fig_count, gallery_conf)
 
@@ -518,7 +533,7 @@ def execute_script(code_block, example_globals, image_path, fig_count,
         os.chdir(cwd)
 
     print(" - time elapsed : %.2g sec" % time_elapsed)
-    code_output = u"\n{0}\n\n{1}\n\n".format(image_list, stdout)
+    code_output = "\n{0}\n\n{1}\n\n".format(image_list, stdout)
 
     return code_output, time_elapsed, fig_count + len(figure_list)
 
@@ -554,7 +569,7 @@ def generate_file_rst(fname, target_dir, src_dir, gallery_conf):
     time_elapsed = 0
 
     ref_fname = example_file.replace(os.path.sep, '_')
-    example_rst = u"""\n\n.. _sphx_glr_{0}:\n\n""".format(ref_fname)
+    example_rst = """\n\n.. _sphx_glr_{0}:\n\n""".format(ref_fname)
     example_nb = Notebook(fname, target_dir)
 
     filename_pattern = gallery_conf.get('filename_pattern')
@@ -613,9 +628,10 @@ def generate_file_rst(fname, target_dir, src_dir, gallery_conf):
 
     time_m, time_s = divmod(time_elapsed, 60)
     example_nb.save_file()
-    with open(os.path.join(target_dir, base_image_name + '.rst'), 'wb') as f:
+    with codecs.open(os.path.join(target_dir, base_image_name + '.rst'),
+                     'w', 'utf-8') as f:
         example_rst += CODE_DOWNLOAD.format(time_m, time_s, fname,
                                             example_nb.file_name)
-        f.write(example_rst.encode('utf-8'))
+        f.write(example_rst)
 
     return amount_of_code, time_elapsed
