Merge pull request #6 from stitchfix/py3

Add Python 3 Support

Author: kzielnicki

README.md

@@ -7,17 +7,18 @@ Nodebook is a plugin for Jupyter Notebook designed to enforce an ordered flow of
 
 ## Installation
 
-Nodebook is available on pypi and can be installed with pip:
+Nodebook is available on pypi and can be installed with pip. Additionally, the jupyter extension must be registered:
 ```
 pip install nodebook
+jupyter nbextension install --py nodebook
 ```
 
 ## Usage
 
 To use Nodebook, add the following lines to a cell in your Jupyter notebook:
 ```
 #pragma nodebook off
-%load_ext nodebookext
+%load_ext nodebook.ipython
 %nodebook {mode} {name}
 ```
 Where `{mode}` is one of `memory` or `disk`, and `{name}` is a unique identifier for your notebook.
@@ -30,9 +31,9 @@ For additional example usage, see [nodebook_demo.ipynb](./nodebook_demo.ipynb).
 
 ## FAQ
 
-#### Q: Does Nodebook support Python 3?
+#### Q: Should I use Python 2 or Python 3 with Nodebook?
 
-Unfortunately, not yet. Please try it in Python 2.7.
+There has been an [increasing consensus](http://www.python3statement.org/) toward sunsetting support for Python 2, including in Project Jupyter. Nodebook currently supports both Python 2 and Python 3, but Python 3 is preferred.
 
 #### Q: Why am I seeing "ERROR:root:Cell magic `%%execute_cell` not found."?
 

nodebook/__init__.py

@@ -0,0 +1,11 @@
+def _jupyter_nbextension_paths():
+    """
+    Jupyter nbextension location
+    """
+    return [dict(
+        section="notebook",
+        src="ipython/nbextensions",
+        # directory in the `nbextension/` namespace
+        dest="nodebook",
+        # _also_ in the `nbextension/` namespace
+        require="nodebook/nodebookext")]

nodebook/ipython/__init__.py

@@ -0,0 +1 @@
+from .nodebookext import load_ipython_extension

ipython/nbextensions/nodebookext.js …book/ipython/nbextensions/nodebookext.jsipython/nbextensions/nodebookext.js renamed to nodebook/ipython/nbextensions/nodebookext.js ipython/nbextensions/nodebookext.js renamed to nodebook/ipython/nbextensions/nodebookext.js

@@ -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
@@ -92,7 +93,7 @@ def execute_cell(line, cell):
 def load_ipython_extension(ipython):
     ipython.register_magic_function(nodebook, magic_kind='line')
     ipython.register_magic_function(execute_cell, magic_kind='cell')
-    ipython.run_cell_magic('javascript', '', "Jupyter.utils.load_extensions('nodebookext')")
+    ipython.run_cell_magic('javascript', '', "Jupyter.utils.load_extensions('nodebook/nodebookext')")
 
 
 def unload_ipython_extension(ipython):

ipython/extensions/nodebookext.py nodebook/ipython/nodebookext.pyipython/extensions/nodebookext.py renamed to nodebook/ipython/nodebookext.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 @@ def run_node(self, node_id):
         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 @@ def _find_latest_output(self, node, var):
         """
         # 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 @@ def _find_latest_output(self, node, var):
             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 @@ def _update_output_hashes(self, node, outputs):
         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 @@ def run(self, input_objs, input_hashes):
         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

nodebook/nodebookcore.py

@@ -1,10 +1,12 @@
+from __future__ import absolute_import
 import io
 import os
 from functools import partial
 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 @@
 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 @@ def __init__(self, persist_path=None):
         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 @@ def get(self, key, default=None):
             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
 

nodebook/pickledict.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 add_dependencies(node_inputs, dep_set):
     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

nodebook/utils.py

@@ -1,3 +1,4 @@
+from __future__ import absolute_import
 from setuptools import setup, find_packages
 import os
 import sys
@@ -8,7 +9,7 @@
 
 setup(
     name='nodebook',
-    version='0.1.0',
+    version='0.2.0',
     author='Kevin Zielnicki',
     author_email='kzielnicki@stitchfix.com',
     license='Stitch Fix 2017',
@@ -17,7 +18,7 @@
     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',
@@ -26,8 +27,7 @@
         'pytest-runner',
     ],
     tests_require=['pytest'],
-    data_files=[
-        (os.path.expanduser('~/.ipython/nbextensions'), ['ipython/nbextensions/nodebookext.js']),
-        (os.path.expanduser('~/.ipython/extensions'), ['ipython/extensions/nodebookext.py']),
-    ],
+    package_data={
+        'nodebook': ['ipython/nbextensions/*.js']
+    },
 )