Supports Python 2

README.md

@@ -3,7 +3,9 @@
 [![PyPI](https://img.shields.io/pypi/v/nb_black.svg)]()
 [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/nb_black.svg)]()
 
-A simple extension for Jupyter Notebook and Jupyter Lab to beautify Python code automatically using **black**.
+A simple extension for Jupyter Notebook and Jupyter Lab to beautify Python code automatically using **Black**.
+
+Please note that since the **Black** package only supports Python 3.6+, so **YAPF** package will be used for the lower versions.
 
 ## Installation
 

nb_black.py

@@ -1,19 +1,38 @@
 #!/usr/bin/python
 # -*- coding: utf-8 -*-
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
 import re
+import sys
+
+if sys.version_info >= (3, 6, 0):
+    from black import format_str
+
+
+    def _format_code(code):
+        return format_str(src_contents=code, line_length=80)
+
+
+else:
+    from yapf.yapflib.yapf_api import FormatCode
+
 
-from black import format_str
+    def _format_code(code):
+        return FormatCode(code, style_config="facebook")
 
 
 class BlackFormatter(object):
     def __init__(self, ip):
         self.shell = ip
 
-    def format(self, result):
-        cell = result.info.raw_cell
+    def format(self):
         try:
+            cell = self.shell.user_ns["In"][-1]
             cell = re.sub(r"^(\s*[!%?])", "# :@BF@: \g<1>", cell, flags=re.M)
-            cell = format_str(src_contents=cell, line_length=88)
+            cell = _format_code(cell)
             cell = re.sub(r"^\s*# :@BF@: (\s*[!%?])", "\g<1>", cell, flags=re.M)
             self.shell.set_next_input(cell.rstrip(), replace=True)
         except ValueError:

setup.py

@@ -7,15 +7,15 @@
 
 def readme(file_name):
     if os.path.isfile(file_name):
-        with open(file_name, "r", encoding="UTF-8") as f:
+        with open(file_name, "r") as f:
             return f.read()
 
 
 setup(
     name="nb_black",
-    version="1.0.0",
+    version="1.0.1",
     description="A simple extension for Jupyter Notebook and Jupyter Lab to beautify Python code automatically using "
-                "black.",
+                "Black.",
     long_description=readme(file_name="README.md"),
     keywords="black-formatter black-beautifier black jupyterlab-extension jupyter-notebook-extension",
     url="https://github.com/dnanhkhoa/nb_black",
@@ -24,11 +24,18 @@ def readme(file_name):
     license="MIT",
     py_modules=["nb_black"],
     zip_safe=False,
-    install_requires=["black"],
+    install_requires=[
+        "yapf; python_version < '3.6'",
+        "black; python_version >= '3.6'",
+        "ipython",
+    ],
     classifiers=[
         "Development Status :: 5 - Production/Stable",
         "Intended Audience :: Developers",
         "License :: OSI Approved :: MIT License",
+        "Programming Language :: Python",
+        "Programming Language :: Python :: 2",
         "Programming Language :: Python :: 3",
+        "Topic :: Software Development :: Libraries :: Python Modules",
     ],
 )