Minor change & add some test cases

Signed-off-by: Xiaochao Dong (@damnever) <the.xcdong@gmail.com>
damnever · Nov 7, 2022 · f8c387d · f8c387d
1 parent 22e8441
commit f8c387d
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 13 deletions.
diff --git a/pigar/parser.py b/pigar/parser.py
@@ -56,22 +56,29 @@ def parse_imports(package_root, ignores=None):
     return imported_modules, user_modules
 
 
-def _filter_ipynb_commands(cell_lines):
-    """
-    Filter out cells with magics and shell commands
-    e.g %matplotlib inline or !pip install pigar
-    """
-    return re.sub(r"^((\s)*(!|%).*$)|(\n\s*%.*)", "", cell_lines)
+# Match ipython notebook magics and shell commands.
+# e.g %matplotlib inline or !pip install pigar
+#
+# Ref:
+#  - https://ipython.readthedocs.io/en/stable/interactive/magics.html
+#  - https://ipython.org/ipython-doc/3/interactive/shell.html
+_ipynb_magics_and_commands_regex = re.compile(
+    r"[^#]*\s*(!|%)[a-zA-Z][a-zA-Z0-9_-]*.*"
+)
 
 
 def _read_code(fpath):
     if fpath.endswith(".ipynb"):
         nb = nbformat.read(fpath, as_version=4)
         code = ""
         for cell in nb.cells:
-            if cell.cell_type == "code":
-                source = _filter_ipynb_commands(cell.source)
-                code += source + "\n"
+            if cell.cell_type != "code":
+                continue
+            for line in cell.source.splitlines():
+                match = _ipynb_magics_and_commands_regex.match(line)
+                if not (match and match.group(0) == line):
+                    code += line
+                code += "\n"
         return code
     elif fpath.endswith(".py"):
         with open(fpath, 'rb') as f:
@@ -96,6 +103,7 @@ def parse_file_imports(fpath, content):
 
 
 class ImportsParser(object):
+
     def __init__(self, rawcode_callback=None):
         self._modules = []
         self._rawcode_callback = rawcode_callback
@@ -308,7 +316,7 @@ def _search_path(path):
                     line = line.strip()
                     if line != '.':
                         dev_dir = line
-            if not dev_dir or not pathlib.exists(dev_dir):
+            if not dev_dir:
                 continue
             # Egg info path.
             info_dir = [

diff --git a/pigar/tests/imports_example/notebook.ipynb b/pigar/tests/imports_example/notebook.ipynb
@@ -23,7 +23,7 @@
       "metadata": {},
       "source": [
         "### Markdown title"
-        ]
+      ]
     },
     {
       "cell_type": "code",
@@ -39,16 +39,29 @@
         }
       ],
       "source": [
-        "import notebook\n",
+        "import notebook # !echo bad\n",
+        "!git status\n",
+        " %cd /tmp\n",
         "print(dir(notebook))"
-        ]
+      ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {},
       "outputs": [],
       "source": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "lines = !ls -l\n",
+        "!echo hello\n",
+        "%alias bracket echo \"knock knock\"\n",
+        "%who_ls\n"
+      ]
     }
   ],
   "metadata": {