Repository: stitchfix/nodebook

Issue:
Python 3 support
Nodebook currently only supports python 2, but python 3 support should be included for the next major release.

Discussion:

Patch:
diff --git a/README.md b/README.md
index ce22d57..5b3672b 100644
--- a/README.md
+++ b/README.md
@@ -30,9 +30,9 @@ For additional example usage, see [nodebook_demo.ipynb](./nodebook_demo.ipynb).
 
 ## FAQ
 
-#### Q: Does Nodebook support Python 3?
+#### Q: Does Nodebook support Python 2?
 
-Unfortunately, not yet. Please try it in Python 2.7.
+Yes, but we recommend using Python 3.
 
 #### Q: Why am I seeing "ERROR:root:Cell magic `%%execute_cell` not found."?
 
diff --git a/ipython/extensions/nodebookext.py b/ipython/extensions/nodebookext.py
index 30dc7a5..f049f7a 100644
--- a/ipython/extensions/nodebookext.py
+++ b/ipython/extensions/nodebookext.py
@@ -1,4 +1,5 @@
-import cPickle as pickle
+from __future__ import absolute_import
+import six.moves.cPickle as pickle
 import os
 import sys
 import errno
diff --git a/nodebook/nodebookcore.py b/nodebook/nodebookcore.py
index a12eda5..98b5cdc 100644
--- a/nodebook/nodebookcore.py
+++ b/nodebook/nodebookcore.py
@@ -1,6 +1,9 @@
+from __future__ import absolute_import
+from __future__ import print_function
 from . import pickledict
 import ast
-import __builtin__
+import six.moves.builtins
+import six
 
 INDENT = '    '  # an indent is canonically 4 spaces ;)
 
@@ -116,7 +119,7 @@ class Nodebook(object):
         res, output_objs, output_hashes = node.run(input_objs, input_hashes)
 
         # update node outputs
-        for var, val in output_objs.iteritems():
+        for var, val in six.iteritems(output_objs):
             self.variables[output_hashes[var]] = val
         self._update_output_hashes(node, output_hashes)
 
@@ -129,7 +132,7 @@ class Nodebook(object):
         """
         # base case
         if node is None:
-            if var in __builtin__.__dict__:
+            if var in six.moves.builtins.__dict__:
                 return None
             else:
                 raise KeyError("name '%s' is not defined" % var)
@@ -141,7 +144,7 @@ class Nodebook(object):
             else:
                 # re-run the parent if it wasn't valid
                 # TODO: synchronize output with frontend javascript
-                print "auto-running invalidated node N_%s (%s)" % (node.get_index() + 1, node.name)
+                print("auto-running invalidated node N_%s (%s)" % (node.get_index() + 1, node.name))
                 self.run_node(node.name)
                 return self._find_latest_output(node, var)
         else:
@@ -153,23 +156,23 @@ class Nodebook(object):
         Update node's output hashes and invalid downstream nodes that depended on their previous values
         """
         # invalidate any any children relying on specific hash-versions of old outputs that aren't in the new outputs
-        invalidated_outputs = set(node.outputs.iteritems()) - set(outputs.iteritems())
+        invalidated_outputs = set(six.iteritems(node.outputs)) - set(six.iteritems(outputs))
         invalidated_outputs = {k: v for k, v in invalidated_outputs}
 
         # also invalidate any children that rely on any version of a brand-new output, regardless of hash
         # TODO this is potentially overly restrictive, if, eg, a value is blindly over-written again later
         # TODO(con't) we should try to account for this to avoid invalidating excessively many cells
-        new_outputs = set(outputs.iteritems()) - set(node.outputs.iteritems())
+        new_outputs = set(six.iteritems(outputs)) - set(six.iteritems(node.outputs))
         new_outputs = {k: v for k, v in new_outputs}
 
         # update reference counts
-        for val_hash in new_outputs.itervalues():
+        for val_hash in six.itervalues(new_outputs):
             self.add_ref(val_hash)
-        for val_hash in invalidated_outputs.itervalues():
+        for val_hash in six.itervalues(invalidated_outputs):
             self.remove_ref(val_hash)
 
         # invalidate changed outputs
-        invalidated_outputs.update({k: None for k, _ in new_outputs.iteritems()})
+        invalidated_outputs.update({k: None for k, _ in six.iteritems(new_outputs)})
         node.outputs = outputs
         node.invalidate_children(invalidated_outputs)
 
@@ -274,10 +277,10 @@ class Node(object):
         block = ast.parse(self.code)
         if len(block.body) > 0 and type(block.body[-1]) is ast.Expr:
             last = ast.Expression(block.body.pop().value)
-            exec compile(block, '<string>', mode='exec') in env
+            exec(compile(block, '<string>', mode='exec'), env)
             res = eval(compile(last, '<string>', mode='eval'), env)
         else:
-            exec compile(block, '<string>', mode='exec') in env
+            exec(compile(block, '<string>', mode='exec'), env)
             res = None
 
         # find outputs which have changed from input hashes
diff --git a/nodebook/pickledict.py b/nodebook/pickledict.py
index 6e8c4be..d4fc6e5 100644
--- a/nodebook/pickledict.py
+++ b/nodebook/pickledict.py
@@ -1,3 +1,4 @@
+from __future__ import absolute_import
 import io
 import os
 from functools import partial
@@ -5,6 +6,7 @@ import hashlib
 import pandas as pd
 import msgpack
 import inspect
+import six
 
 # using dill instead of pickle for more complete serialization
 import dill
@@ -13,9 +15,18 @@ import dill
 try:
     from cStringIO import StringIO
 except ImportError:
-    from StringIO import StringIO
+    try:
+        from StringIO import StringIO
+    except ImportError:
+        # Python3. We are using StringIO as a target for pickle, so we
+        # actually want BytesIO.
+        from io import BytesIO as StringIO
 
-import UserDict
+try:
+    from UserDict import DictMixin
+except ImportError:
+    # see https://github.com/flask-restful/flask-restful/pull/231/files
+    from collections import MutableMapping as DictMixin
 
 PANDAS_CODE = 1
 DILL_CODE = 2
@@ -45,7 +56,7 @@ def msgpack_deserialize(code, data):
         return msgpack.ExtType(code, data)
 
 
-class PickleDict(object, UserDict.DictMixin):
+class PickleDict(DictMixin):
     """
     Dictionary with immutable elements using pickle(dill), optionally supporting persisting to disk
     """
@@ -55,13 +66,13 @@ class PickleDict(object, UserDict.DictMixin):
         persist_path: if provided, perform serialization to/from disk to this path
         """
         self.persist_path = persist_path
+        self.encodings = {}
         self.dump = partial(msgpack.dump, default=msgpack_serialize)
         self.load = partial(msgpack.load, ext_hook=msgpack_deserialize)
-
         self.dict = {}
 
     def keys(self):
-        return self.dict.keys()
+        return list(self.dict.keys())
 
     def __len__(self):
         return len(self.dict)
@@ -77,25 +88,33 @@ class PickleDict(object, UserDict.DictMixin):
             return self[key]
         return default
 
+    def __iter__(self):
+        for key in self.dict:
+            yield key
+
     def __getitem__(self, key):
         if self.persist_path is not None:
             path = self.dict[key]
             with open(path, 'rb') as f:
-                value = self.load(f)
+                value = self.load(f, encoding=self.encodings[key])
         else:
             f = StringIO(self.dict[key])
-            value = self.load(f)
+            value = self.load(f, encoding=self.encodings[key])
         return value
 
     def __setitem__(self, key, value):
+        encoding = None
+        if isinstance(value, six.string_types):
+            encoding = 'utf-8'
+        self.encodings[key] = encoding
         if self.persist_path is not None:
             path = os.path.join(self.persist_path, '%s.pak' % key)
             with open(path, 'wb') as f:
-                self.dump(value, f)
+                self.dump(value, f, encoding=encoding)
             self.dict[key] = path
         else:
             f = StringIO()
-            self.dump(value, f)
+            self.dump(value, f, encoding=encoding)
             serialized = f.getvalue()
             self.dict[key] = serialized
 
diff --git a/nodebook/utils.py b/nodebook/utils.py
index 5bffb80..baab14e 100644
--- a/nodebook/utils.py
+++ b/nodebook/utils.py
@@ -1,10 +1,12 @@
+from __future__ import absolute_import
 import json
-from nodebookcore import INDENT
+from .nodebookcore import INDENT
+import six
 
 
 def output_to_function(output_node, main_closing_statement, args):
     def add_dependencies(node_inputs, dep_set):
-        dep = {(k, v) for k, v in node_inputs.iteritems() if k not in args and v is not None}
+        dep = {(k, v) for k, v in six.iteritems(node_inputs) if k not in args and v is not None}
         return dep_set.union(dep)
 
     depends = add_dependencies(output_node.inputs, set())
@@ -14,8 +16,8 @@ def output_to_function(output_node, main_closing_statement, args):
     n = output_node
     while not depends.issubset(avail) and n.parent is not None:
         n = n.parent
-        if len(depends.intersection(n.outputs.iteritems())) != 0:
-            avail.update(n.outputs.iteritems())
+        if len(depends.intersection(six.iteritems(n.outputs))) != 0:
+            avail.update(six.iteritems(n.outputs))
             depends = add_dependencies(n.inputs, depends)
             funcs.append(n.extract_function())
 
@@ -37,16 +39,16 @@ def create_module(node, export_statement, input_dict):
     """
     create a python module to execute a given node
     """
-    body = output_to_function(node, export_statement, input_dict.keys())
+    body = output_to_function(node, export_statement, list(input_dict.keys()))
 
     imports = '\n'.join([
         'import json',
     ])
 
     deser = ''
-    for k, v in input_dict.iteritems():
+    for k, v in six.iteritems(input_dict):
         deser += "\n{} = json.loads('{}')".format(k, json.dumps(v))
-    deser += '\nmain({})'.format(','.join(input_dict.iterkeys()))
+    deser += '\nmain({})'.format(','.join(six.iterkeys(input_dict)))
 
     code = '{}\n\n{}\n\n{}\n'.format(imports, body, deser)
     return code
diff --git a/setup.py b/setup.py
index b47742a..6a5b93a 100644
--- a/setup.py
+++ b/setup.py
@@ -1,3 +1,4 @@
+from __future__ import absolute_import
 from setuptools import setup, find_packages
 import os
 import sys
@@ -8,7 +9,7 @@ if 'bdist_wheel' in sys.argv:
 
 setup(
     name='nodebook',
-    version='0.1.0',
+    version='0.2.0-dev',
     author='Kevin Zielnicki',
     author_email='kzielnicki@stitchfix.com',
     license='Stitch Fix 2017',
@@ -17,7 +18,7 @@ setup(
     long_description='Nodebook Jupyter Extension',
     url='https://github.com/stitchfix/nodebook',
     install_requires=[
-        'ipython<6',  # newer versions of ipython do not support 2.7
+        'ipython',
         'jupyter',
         'click',
         'dill',


Test patch:
diff --git a/tests/test_nodebookcore.py b/tests/test_nodebookcore.py
index dcb2579..cfd9646 100644
--- a/tests/test_nodebookcore.py
+++ b/tests/test_nodebookcore.py
@@ -1,3 +1,4 @@
+from __future__ import absolute_import
 import pandas as pd
 import pytest
 from nodebook.nodebookcore import ReferenceFinder, Nodebook, Node
diff --git a/tests/test_pickledict.py b/tests/test_pickledict.py
index b6d773c..ef35fdd 100644
--- a/tests/test_pickledict.py
+++ b/tests/test_pickledict.py
@@ -1,3 +1,6 @@
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import unicode_literals
 import pandas as pd
 import pytest
 from nodebook.pickledict import PickleDict
@@ -7,7 +10,7 @@ from nodebook.pickledict import PickleDict
 def mydict(request, tmpdir):
     if request.param == 'tmpdir':
         persist_path = tmpdir.strpath
-        print persist_path
+        print(persist_path)
     else:
         persist_path = None
     return PickleDict(persist_path=persist_path)
@@ -22,6 +25,10 @@ class TestPickleDict(object):
         mydict['test_string'] = 'foo'
         assert mydict['test_string'] == 'foo'
 
+    def test_bytes(self, mydict):
+        mydict['test_bytes'] = b'foo'
+        assert mydict['test_bytes'] == b'foo'
+
     def test_df(self, mydict):
         df = pd.DataFrame({'a': [0, 1, 2], 'b': ['foo', 'bar', 'baz']})
         mydict['test_df'] = df