Add simple implementation of Python 3 style open()

takluyver · takluyver · commit d60f657e2b75 · 2011-09-07T12:18:47.000+01:00
diff --git a/IPython/core/magic.py b/IPython/core/magic.py
@@ -52,6 +52,7 @@
 from IPython.core.prefilter import ESC_MAGIC
 from IPython.lib.pylabtools import mpl_runner
 from IPython.testing.skipdoctest import skip_doctest
+from IPython.utils import py3compat
 from IPython.utils.io import file_read, nlprint
 from IPython.utils.path import get_py_filename, unquote_filename
 from IPython.utils.process import arg_split, abbrev_cwd
@@ -2085,11 +2086,9 @@ def magic_save(self,parameter_s = ''):
         except (TypeError, ValueError) as e:
             print e.args[0]
             return
-        if isinstance(cmds, unicode):
-            cmds = cmds.encode("utf-8")
-        with open(fname,'w') as f:
-            f.write("# coding: utf-8\n")
-            f.write(cmds)
+        with py3compat.open(fname,'w', encoding="utf-8") as f:
+            f.write(u"# coding: utf-8\n")
+            f.write(py3compat.cast_unicode(cmds))
         print 'The following commands were written to file `%s`:' % fname
         print cmds
     
diff --git a/IPython/utils/py3compat.py b/IPython/utils/py3compat.py
@@ -2,6 +2,8 @@
 """Compatibility tricks for Python 3. Mainly to do with unicode."""
 import sys
 
+orig_open = open
+
 def no_code(x, encoding=None):
     return x
 
@@ -38,6 +40,8 @@ def isidentifier(s, dotted=False):
         if dotted:
             return all(isidentifier(a) for a in s.split("."))
         return s.isidentifier()
+    
+    open = orig_open
 
 else:
     PY3 = False
@@ -56,6 +60,27 @@ def isidentifier(s, dotted=False):
         if dotted:
             return all(isidentifier(a) for a in s.split("."))
         return bool(_name_re.match(s))
+    
+    class open(object):
+        """Wrapper providing key part of Python 3 open() interface."""
+        def __init__(self, fname, mode="r", encoding="utf-8"):
+            self.f = orig_open(fname, mode)
+            self.enc = encoding
+        
+        def write(self, s):
+            return self.f.write(s.encode(self.enc))
+        
+        def read(self, size=-1):
+            return self.f.read(size).decode(self.enc)
+        
+        def close(self):
+            return self.f.close()
+        
+        def __enter__(self):
+            return self
+        
+        def __exit__(self, etype, value, traceback):
+            self.f.close()
 
 def execfile(fname, glob, loc=None):
     loc = loc if (loc is not None) else glob
