Merge pull request #54 from asmeurer/optional-nb

Make notebook support and dependencies optional

Author: Saransh-cpp

.github/workflows/ci.yml

@@ -37,11 +37,19 @@ jobs:
           python-version: ${{ matrix.python-version }}
           allow-prereleases: true
 
-      - name: Install removestar
+      - name: Install removestar (lite)
         run: |
           python -m pip install -e ".[dev]"
 
-      - name: Test removestar
+      - name: Test removestar (lite)
+        run: |
+          pytest --cov=removestar . -vv
+
+      - name: Install removestar (nb)
+        run: |
+          python -m pip install -e ".[dev,nb]"
+
+      - name: Test removestar (nb)
         run: |
           pytest --cov=removestar . -vv
 

README.md

@@ -22,6 +22,7 @@ Install `removestar` globally to use it through CLI using `pypi` -
 
 ```bash
 pip install removestar
+pip install "removestar[nb]"  # notebook support
 ```
 
 or `conda` -
@@ -39,7 +40,8 @@ or add `removestar` in `.pre-commit-config.yaml` -
     - id: removestar
       args: [-i] # See docs for all args (-i edits file in-place)
       additional_dependencies: # The libraries or packages your code imports
-        - ... # Should be . if running inside a library (to install the library itself in the environment)
+        - ... # Add . if running inside a library (to install the library itself in the environment)
+        - ... # Add nbformat and nbconvert for notebook support
 ```
 
 ## Usage
@@ -72,7 +74,7 @@ $ removestar -i file.py # Edits file.py in-place
 
 $ removestar -i module/ # Modifies every Python file in module/ recursively
 
-# notebooks
+# notebooks (make sure nbformat and nbconvert are installed)
 
 $ removestar file.ipynb # Shows diff but does not edit file.ipynb
 

pyproject.toml

@@ -10,30 +10,33 @@ description = "A tool to automatically replace 'import *' imports with explicit
 license = "MIT"
 requires-python = ">=3.7"
 authors = [
-    { name = "Aaron Meurer", email = "[email protected]" },
+  { name = "Aaron Meurer", email = "[email protected]" },
 ]
 maintainers = [
-    { name = "Aaron Meurer", email = "[email protected]" },
-    { name = "Saransh Chopra", email = "[email protected]" }
+  { name = "Aaron Meurer", email = "[email protected]" },
+  { name = "Saransh Chopra", email = "[email protected]" }
 ]
 classifiers = [
-    "Development Status :: 5 - Production/Stable",
-    "License :: OSI Approved :: MIT License",
-    "Operating System :: OS Independent",
-    "Programming Language :: Python",
-    "Programming Language :: Python :: 3 :: Only",
-    "Programming Language :: Python :: 3.8",
-    "Programming Language :: Python :: 3.9",
-    "Programming Language :: Python :: 3.10",
-    "Programming Language :: Python :: 3.11",
+  "Development Status :: 5 - Production/Stable",
+  "License :: OSI Approved :: MIT License",
+  "Operating System :: OS Independent",
+  "Programming Language :: Python",
+  "Programming Language :: Python :: 3 :: Only",
+  "Programming Language :: Python :: 3.8",
+  "Programming Language :: Python :: 3.9",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
 ]
 dependencies = [
-    "pyflakes",
-    "nbformat",
-    "nbconvert",
+  "pyflakes",
 ]
 
 [project.optional-dependencies]
+nb = [
+  "nbformat",
+  "nbconvert",
+]
 dev = [
   "pytest>=6",
   "pytest-cov>=3",

removestar/__main__.py

@@ -16,18 +16,16 @@
 
 import argparse
 import glob
+import importlib.util
 import io
 import os
 import sys
 import tempfile
 
-import nbformat
-from nbconvert import PythonExporter
-
 from . import __version__
 from .helper import get_diff_text
 from .output import get_colored_diff, red
-from .removestar import fix_code, replace_in_nb
+from .removestar import fix_code
 
 
 class RawDescriptionHelpArgumentDefaultsHelpFormatter(
@@ -92,6 +90,14 @@ def main():  # noqa: PLR0912, C901
     if args.max_line_length == 0:
         args.max_line_length = float("inf")
 
+    try:
+        import nbformat
+        from nbconvert import PythonExporter
+
+        from .removestar import replace_in_nb
+    except ImportError:
+        pass
+
     exit_1 = False
     for file in _iter_paths(args.paths):
         _, filename = os.path.split(file)
@@ -144,7 +150,11 @@ def main():  # noqa: PLR0912, C901
                             )
                         )
                     )
-        elif file.endswith(".ipynb"):
+        elif (
+            file.endswith(".ipynb")
+            and importlib.util.find_spec("nbconvert") is not None
+            and importlib.util.find_spec("nbformat") is not None
+        ):
             tmp_file = tempfile.NamedTemporaryFile()
             tmp_path = tmp_file.name
 

removestar/removestar.py

@@ -1,15 +1,18 @@
 import ast
 import builtins
+import contextlib
 import os
 import re
 import sys
 from functools import lru_cache
 from pathlib import Path
 
-from nbconvert import NotebookExporter
 from pyflakes.checker import _MAGIC_GLOBALS, Checker, ModuleScope
 from pyflakes.messages import ImportStarUsage, ImportStarUsed
 
+with contextlib.suppress(ImportError):
+    from nbconvert import NotebookExporter
+
 from .output import green, yellow
 
 # quit and exit are not included in old versions of pyflakes

tests/test_removestar_nb.py

@@ -1,11 +1,13 @@
 import os
 import tempfile
 
-import nbformat as nbf
-from nbconvert import NotebookExporter, PythonExporter
+import pytest
 
 from removestar.removestar import fix_code, replace_in_nb
 
+nbf = pytest.importorskip("nbformat")
+nbc = pytest.importorskip("nbconvert")
+
 fixed_code = """#!/usr/bin/env python
 # coding: utf-8
 # # Notebook for testing.
@@ -35,7 +37,7 @@ def prepare_nb(output_path="_test.ipynb"):
     with open(output_path) as f:
         nb = nbf.reads(f.read(), nbf.NO_CONVERT)
 
-    exporter = PythonExporter()
+    exporter = nbc.PythonExporter()
     code, _ = exporter.from_notebook_node(nb)
     tmp_file.write(code.encode("utf-8"))
 
@@ -78,7 +80,7 @@ def test_replace_nb():
     with open("_test.ipynb", "w+") as f:
         f.writelines(fixed_code)
 
-    exporter = NotebookExporter()
+    exporter = nbc.NotebookExporter()
     code, _ = exporter.from_filename("_test.ipynb")
 
     assert code == fixed_code