wip

pipreqs/pipreqs.py

@@ -49,20 +49,15 @@
 from yarg import json2package
 from yarg.exceptions import HTTPError
 
-try:
-    PythonExporter = None
-    ignore_notebooks = False
-    from nbconvert import PythonExporter
-except ImportError:
-    pass
-
 from pipreqs import __version__
 
 REGEXP = [
     re.compile(r"^import (.+)$"),
     re.compile(r"^from ((?!\.+).*?) import (?:.*)$"),
 ]
 
+ignore_notebooks = False
+
 
 @contextmanager
 def _open(filename=None, mode="r"):
@@ -117,22 +112,20 @@ def get_all_imports(path, encoding="utf-8", extra_ignore_dirs=None, follow_links
             ignore_dirs_parsed.append(os.path.basename(os.path.realpath(e)))
         ignore_dirs.extend(ignore_dirs_parsed)
 
+    def get_file_extensions():
+        return [".py"] if ignore_notebooks else [".py", ".ipynb"]
+
+    extensions = get_file_extensions()
+
     walk = os.walk(path, followlinks=follow_links)
     for root, dirs, files in walk:
         dirs[:] = [d for d in dirs if d not in ignore_dirs]
 
         candidates.append(os.path.basename(root))
-        if notebooks_are_enabled():
-            files = [fn for fn in files if file_ext_is_allowed(fn, [".py", ".ipynb"])]
-        else:
-            files = [fn for fn in files if file_ext_is_allowed(fn, [".py"])]
+        py_files = [file for file in files if file_ext_is_allowed(file, [".py"])]
+        candidates.extend([os.path.splitext(filename)[0] for filename in py_files])
 
-        candidates = list(
-            map(
-                lambda fn: os.path.splitext(fn)[0],
-                filter(lambda fn: file_ext_is_allowed(fn, [".py"]), files),
-            )
-        )
+        files = [fn for fn in files if file_ext_is_allowed(fn, extensions)]
 
         for file_name in files:
             file_name = os.path.join(root, file_name)
@@ -178,10 +171,6 @@ def get_all_imports(path, encoding="utf-8", extra_ignore_dirs=None, follow_links
     return list(packages - data)
 
 
-def notebooks_are_enabled():
-    return PythonExporter and not ignore_notebooks
-
-
 def file_ext_is_allowed(file_name, acceptable):
     return os.path.splitext(file_name)[1] in acceptable
 
@@ -197,7 +186,7 @@ def ipynb_2_py(file_name, encoding="utf-8"):
         str: parsed string
 
     """
-
+    
     exporter = PythonExporter()
     (body, _) = exporter.from_filename(file_name)
 
@@ -493,7 +482,9 @@ def init(args):
     encoding = args.get("--encoding")
     extra_ignore_dirs = args.get("--ignore")
     follow_links = not args.get("--no-follow-links")
-    ignore_notebooks = args.get("--ignore-notebooks")
+
+    ignore_notebooks = args.get("--ignore-notebooks", False)
+
     input_path = args["<path>"]
 
     if encoding is None: