Repository: microsoft/pybryt

Issue:
Import annotation
Create an annotation that asserts that a library was imported. This is needed because Value annotations can't serialize modules when they're saved.

e.g.

```python
import numpy as np
pybryt.Import(np)
```

Discussion:

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2f004e5..aaaba8f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,7 @@ All notable changes to this project will be documented in this file, and this pr
 
 ## Unreleased
 
+* Added the `RequireImport` and `ForbidImport` annotations per [#123](https://github.com/microsoft/pybryt/issues/123)
 * Added `TimeComplexityChecker` for analyzing complexity without annotations per [#137](https://github.com/microsoft/pybryt/issues/137)
 
 ## 0.3.1 - 2021-12-01
diff --git a/Makefile b/Makefile
index b0e34d2..b8a0960 100644
--- a/Makefile
+++ b/Makefile
@@ -36,7 +36,7 @@ testcov:
 	$(COVERAGE) run -m pytest $(TESTPATH) $(PYTESTOPTS) 
 	$(COVERAGE) combine
 
-covhtml: testcov
+htmlcov: testcov
 	$(COVERAGE) html
 
 .PHONY: docs
diff --git a/SUPPORT.md b/SUPPORT.md
index dc72f0e..c837b0a 100644
--- a/SUPPORT.md
+++ b/SUPPORT.md
@@ -1,25 +1,13 @@
-# TODO: The maintainer of this repo has not yet edited this file
-
-**REPO OWNER**: Do you want Customer Service & Support (CSS) support for this product/project?
-
-- **No CSS support:** Fill out this template with information about how to file issues and get help.
-- **Yes CSS support:** Fill out an intake form at [aka.ms/spot](https://aka.ms/spot). CSS will work with/help you to determine next steps. More details also available at [aka.ms/onboardsupport](https://aka.ms/onboardsupport).
-- **Not sure?** Fill out a SPOT intake as though the answer were "Yes". CSS will help you decide.
-
-*Then remove this first heading from this SUPPORT.MD file before publishing your repo.*
-
 # Support
 
 ## How to file issues and get help  
 
 This project uses GitHub Issues to track bugs and feature requests. Please search the existing 
-issues before filing new issues to avoid duplicates.  For new issues, file your bug or 
+issues before filing new issues to avoid duplicates. For new issues, file your bug or 
 feature request as a new Issue.
 
-For help and questions about using this project, please **REPO MAINTAINER: INSERT INSTRUCTIONS HERE 
-FOR HOW TO ENGAGE REPO OWNERS OR COMMUNITY FOR HELP. COULD BE A STACK OVERFLOW TAG OR OTHER
-CHANNEL. WHERE WILL YOU HELP PEOPLE?**.
+For help and questions about using this project, please open an issue.
 
 ## Microsoft Support Policy  
 
-Support for this **PROJECT or PRODUCT** is limited to the resources listed above.
+Support for this project is limited to the resources listed above.
diff --git a/docs/annotations/import_annotations.rst b/docs/annotations/import_annotations.rst
new file mode 100644
index 0000000..484d1d4
--- /dev/null
+++ b/docs/annotations/import_annotations.rst
@@ -0,0 +1,40 @@
+.. _import:
+
+Import Annotations
+==================
+
+Import annotations can be used to either require or forbid the use of specific libraries. When PyBryt
+executes a notebook, it parses the AST of each code cell to determine any libraries imported in the
+code. It also captures the modules that any functions or classes it finds belong to when it is
+tracing through the notebook's code. These are collected into a set of library names, which are 
+included in the memory footprint.
+
+All import annotations are subclasses of the abstract 
+:py:class:`ImportAnnotation<pybryt.annotations.import_.ImportAnnotation>` class. When these
+annotations are instantiated, they attempt to import the module to ensure that the string passed
+for the module name is valid. This means that any modules used in import annotations must be installed
+in the environment in which the reference implementation is compiled.
+
+
+Require Import
+--------------
+
+To require an import, use the :py:class:`RequireImport<pybryt.annotations.import_.RequireImport>`
+annotation, which takes the name of the library as a string (along with the usual set of arguments
+that annotation constructors accept).
+
+.. code-block:: python
+
+    pybryt.RequireImport("pandas")
+
+
+Forbid Import
+-------------
+
+To forbid an import, use the :py:class:`ForbidImport<pybryt.annotations.import_.ForbidImport>`
+annotation, which, like the :py:class:`RequireImport<pybryt.annotations.import_.RequireImport>`
+annotation, takes the name of the libary as its only positional argument.
+
+.. code-block:: python
+
+    pybryt.ForbidImport("numpy")
diff --git a/docs/annotations/index.rst b/docs/annotations/index.rst
index b12d022..7533075 100644
--- a/docs/annotations/index.rst
+++ b/docs/annotations/index.rst
@@ -11,6 +11,7 @@ Annotations
     relational_annotations
     complexity_annotations
     type_annotations
+    import_annotations
     collections
     invariants
 
@@ -31,6 +32,7 @@ are five main types of annotations:
 * :ref:`relational annotations<relational>`
 * :ref:`complexity annotations<complexity>`
 * :ref:`type annotations<type>`
+* :ref:`import annotations<import>`
 * :ref:`annotation collections<collections>`
 
 
diff --git a/docs/api_reference.rst b/docs/api_reference.rst
index 5ba3c80..2fabae5 100644
--- a/docs/api_reference.rst
+++ b/docs/api_reference.rst
@@ -64,6 +64,14 @@ Type Annotations
    :undoc-members:
 
 
+Import Annotations
+++++++++++++++++++
+
+.. automodule:: pybryt.annotations.import_
+   :members:
+   :undoc-members:
+
+
 Annotation Collections
 ++++++++++++++++++++++
 
diff --git a/pybryt/annotations/__init__.py b/pybryt/annotations/__init__.py
index 1db514b..29cbe45 100644
--- a/pybryt/annotations/__init__.py
+++ b/pybryt/annotations/__init__.py
@@ -4,6 +4,7 @@ from . import invariants
 from .annotation import *
 from .collection import *
 from .complexity import *
+from .import_ import *
 from .relation import *
 from .type_ import *
 from .value import *
diff --git a/pybryt/annotations/import_.py b/pybryt/annotations/import_.py
new file mode 100644
index 0000000..6da0091
--- /dev/null
+++ b/pybryt/annotations/import_.py
@@ -0,0 +1,134 @@
+"""Annotations for checking imported modules within memory footprints"""
+
+__all__ = ["ForbidImport", "ImportAnnotation", "RequireImport"]
+
+import importlib
+import sys
+
+from abc import abstractmethod
+from typing import Any, Dict, List
+
+from .annotation import Annotation, AnnotationResult
+
+from ..execution import MemoryFootprint
+
+
+class ImportAnnotation(Annotation):
+    """
+    Annotation class for asserting conditions on the set of imported modules
+
+    The set of modules imported in the student implementation is determined using ``sys.modules``,
+    meaning that subclasses count any modules imported by third-party libaries, not just directly
+    by the student's code.
+
+    Args:
+        module (``str``): the module name to forbid in importable form
+        **kwargs: additional keyword arguments passed to the 
+            :py:class:`Annotation<pybryt.annotations.annotation.Annotation>` constructor
+    """
+
+    module: str
+    """the module name to forbid"""
+
+    def __init__(self, module, **kwargs):
+        if not isinstance(module, str):
+            raise TypeError(f"{module} is not a string")
+
+        try:
+            already_imported = module in sys.modules
+            importlib.import_module(module)
+            if not already_imported:  # clean up sys.modules if needed
+                sys.modules.pop(module)
+        except Exception as e:
+            raise ValueError(f"{module} is not importable")
+
+        self.module = module
+
+        super().__init__(**kwargs)
+
+    @property
+    def children(self) -> List[Annotation]:
+        return []
+
+    @abstractmethod
+    def check(self, footprint: MemoryFootprint) -> AnnotationResult:
+        ... # pragma: no cover
+
+    def __eq__(self, other: Any) -> bool:
+        """
+        Checks whether this annotation is equal to another object.
+
+        To be equal to a ``ForbidImport``, the other object must also be a ``ForbidImport`` object, 
+        have the same instance variables, and the same forbidden module.
+
+        Args:
+            other (``object``): the object to compare to
+
+        Returns:
+            ``bool``: whether the objects are equal
+        """
+        return super().__eq__(other) and self.module == other.module
+
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Converts this annotation's details to a JSON-friendly dictionary format.
+
+        Output dictionary contains the annotation's name, group, limit number, success message, and
+        failure message, as well as the module name.
+
+        Returns:
+            ``dict[str, object]``: the dictionary representation of this annotation
+        """
+        d = super().to_dict()
+        d.update({
+            "module": str(self.module),
+        })
+        return d
+
+
+class RequireImport(ImportAnnotation):
+    """
+    Annotation class for asserting that a specific module was imported.
+
+    Args:
+        module (``str``): the module name to forbid in importable form
+        **kwargs: additional keyword arguments passed to the 
+            :py:class:`Annotation<pybryt.annotations.annotation.Annotation>` constructor
+    """
+
+    def check(self, footprint: MemoryFootprint) -> AnnotationResult:
+        """
+        Checks that the memory footprint's imports contain the specified module.
+
+        Args:
+            footprint (:py:class:`pybryt.execution.memory_footprint.MemoryFootprint`): the
+                memory footprint to check against
+        
+        Returns:
+            :py:class:`AnnotationResult`: the results of this annotation against ``footprint``
+        """
+        return AnnotationResult(self.module in footprint.imports, self)
+
+
+class ForbidImport(ImportAnnotation):
+    """
+    Annotation class for asserting that a specific module was not imported.
+
+    Args:
+        module (``str``): the module name to forbid in importable form
+        **kwargs: additional keyword arguments passed to the 
+            :py:class:`Annotation<pybryt.annotations.annotation.Annotation>` constructor
+    """
+
+    def check(self, footprint: MemoryFootprint) -> AnnotationResult:
+        """
+        Checks that the memory footprint's imports don't contain the specified module.
+
+        Args:
+            footprint (:py:class:`pybryt.execution.memory_footprint.MemoryFootprint`): the
+                memory footprint to check against
+        
+        Returns:
+            :py:class:`AnnotationResult`: the results of this annotation against ``footprint``
+        """
+        return AnnotationResult(self.module not in footprint.imports, self)
diff --git a/pybryt/execution/__init__.py b/pybryt/execution/__init__.py
index 2d3838f..678a731 100644
--- a/pybryt/execution/__init__.py
+++ b/pybryt/execution/__init__.py
@@ -24,7 +24,7 @@ from .tracing import (
     TRACING_VARNAME,
 )
 
-from ..preprocessors import IntermediateVariablePreprocessor
+from ..preprocessors import NotebookPreprocessor
 from ..utils import make_secret
 
 
@@ -58,7 +58,7 @@ def execute_notebook(
         :py:class:`pybryt.execution.memory_footprint.MemoryFootprint`: the memory footprint
     """
     nb = deepcopy(nb)
-    preprocessor = IntermediateVariablePreprocessor()
+    preprocessor = NotebookPreprocessor()
     nb = preprocessor.preprocess(nb)
 
     secret = make_secret()
@@ -78,8 +78,6 @@ def execute_notebook(
     last_cell = nbformat.v4.new_code_cell(dedent(f"""\
         from pybryt.execution import tracing_off
         tracing_off()
-        import sys
-        {footprint_varname}.add_imports(*sys.modules.keys())
         {footprint_varname}.filter_out_unpicklable_values()
         import dill
         with open("{footprint_fp}", "wb+") as f:
@@ -94,10 +92,11 @@ def execute_notebook(
     ep.preprocess(nb)
 
     with open(footprint_fp, "rb") as f:
-        footprint = dill.load(f)
+        footprint: MemoryFootprint = dill.load(f)
 
     os.remove(footprint_fp)
 
+    footprint.add_imports(*preprocessor.get_imports())
     footprint.set_executed_notebook(nb)
 
     return footprint
diff --git a/pybryt/execution/tracing.py b/pybryt/execution/tracing.py
index 102ce61..56d1cb4 100644
--- a/pybryt/execution/tracing.py
+++ b/pybryt/execution/tracing.py
@@ -21,8 +21,8 @@ TRACING_VARNAME = "__PYBRYT_TRACING__"
 
 
 def create_collector(
-    skip_types: List[type] = [type, type(len), ModuleType, FunctionType], 
-    addl_filenames: List[str] = []
+    skip_types: List[type] = [type, type(len), FunctionType],
+    addl_filenames: List[str] = [],
 ) -> Tuple[MemoryFootprint, Callable[[FrameType, str, Any], Callable]]:
     """
     Creates a memory footprint to collect observed values and a trace function.
@@ -56,9 +56,16 @@ def create_collector(
             seen_at (``int``, optional): an overriding step counter value
         """
         try:
+            if hasattr(val, "__module__"):
+                footprint.add_imports(val.__module__.split(".")[0])
+
             if type(val) in skip_types:
                 return
 
+            if isinstance(val, ModuleType):
+                footprint.add_imports(val.__name__.split(".")[0])
+                return
+
             footprint.add_value(copy(val), seen_at)
 
         # if something fails, don't track
diff --git a/pybryt/preprocessors/__init__.py b/pybryt/preprocessors/__init__.py
index 97bb945..4686baf 100644
--- a/pybryt/preprocessors/__init__.py
+++ b/pybryt/preprocessors/__init__.py
@@ -1,3 +1,44 @@
 """Submission preprocessors for PyBryt"""
 
+import nbformat
+
+from typing import List, Set
+
+from .abstract_preprocessor import AbstractPreprocessor
+from .imports import ImportFindingPreprocessor
 from .intermediate_variables import IntermediateVariablePreprocessor
+
+
+PREPROCESSORS = [
+    ImportFindingPreprocessor,
+    IntermediateVariablePreprocessor
+]
+
+
+class NotebookPreprocessor(AbstractPreprocessor):
+    """
+    A class for applying a series of preprocessors to a notebook.
+    """
+
+    preprocessors: List[AbstractPreprocessor]
+
+    def __init__(self) -> None:
+        self.preprocessors = []
+        for preprocessor_class in PREPROCESSORS:
+            self.preprocessors.append(preprocessor_class())
+
+    def preprocess(self, nb: nbformat.NotebookNode) -> nbformat.NotebookNode:
+        for preprocessor in self.preprocessors:
+            nb = preprocessor.preprocess(nb)
+
+        return nb
+
+    def get_imports(self) -> Set[str]:
+        """
+        Get the set of imports from the :py:class:`ImportFinder<pybryt.preprocessors.imports.ImportFindingPreprocessor`.
+
+        Returns:
+            ``set[str]``: the set of modules imported
+        """
+        import_finder_idx = [isinstance(p, ImportFindingPreprocessor) for p in self.preprocessors].index(True)
+        return self.preprocessors[import_finder_idx].imports
diff --git a/pybryt/preprocessors/abstract_preprocessor.py b/pybryt/preprocessors/abstract_preprocessor.py
new file mode 100644
index 0000000..b7e4a29
--- /dev/null
+++ b/pybryt/preprocessors/abstract_preprocessor.py
@@ -0,0 +1,30 @@
+"""Abstract base class for notebook preprocessors"""
+
+import nbformat
+
+from abc import ABC, abstractmethod
+from IPython.core.inputtransformer2 import TransformerManager
+
+
+class AbstractPreprocessor(ABC):
+    """
+    An abstract base class for notebook execution preprocessors.
+    """
+
+    transformer_manager: TransformerManager
+
+    def __init__(self) -> None:
+        self.transformer_manager = TransformerManager()
+
+    @abstractmethod
+    def preprocess(self, nb: nbformat.NotebookNode) -> nbformat.NotebookNode:
+        """
+        Preprocesses a notebook for execution.
+
+        Args:
+            nb (``nbformat.NotebookNode``): the notebook to be preprocessed
+
+        Returns:
+            ``nbformat.NotebookNode``: the updated notebook
+        """
+        ...  # pragma: no cover
diff --git a/pybryt/preprocessors/imports.py b/pybryt/preprocessors/imports.py
new file mode 100644
index 0000000..53394f9
--- /dev/null
+++ b/pybryt/preprocessors/imports.py
@@ -0,0 +1,63 @@
+"""Preprocessor for collecting the set of imported modules"""
+
+import ast
+import nbformat
+
+from typing import Set
+
+from .abstract_preprocessor import AbstractPreprocessor
+
+
+class ImportFinder(ast.NodeVisitor):
+    """
+    An AST walker for collecting imported modules.
+    """
+
+    def __init__(self) -> None:
+        self.imports = set()
+
+    def visit_Import(self, node: ast.AST) -> None:
+        """
+        Visit an import statement and collect the names of the imported modules
+        """
+        for alias in node.names:
+            self.imports.add(alias.name.split(".")[0])
+
+    def visit_ImportFrom(self, node: ast.AST) -> None:
+        """
+        Visit an import from statement and collect the names of the imported modules
+        """
+        self.imports.add(node.module)
+
+
+class ImportFindingPreprocessor(AbstractPreprocessor):
+    """
+    A notebook preprocessor that collects the set of modules import in the notebook.
+    """
+
+    imports: Set[str]
+
+    def __init__(self) -> None:
+        self.imports = set()
+        super().__init__()
+
+    def preprocess(self, nb: nbformat.NotebookNode) -> nbformat.NotebookNode:
+        """
+        Populate ``self.imports`` by walking the AST of each code cell.
+
+        Args:
+            nb (``nbformat.NotebookNode``): the notebook to be preprocessed
+
+        Returns:
+            ``nbformat.NotebookNode``: the notebook, unchanged
+        """
+        for cell in nb['cells']:
+            if cell['cell_type'] == 'code':
+                code = cell['source']
+                code = self.transformer_manager.transform_cell(code)
+                tree = ast.parse(code)
+                import_finder = ImportFinder()
+                import_finder.visit(tree)
+                self.imports.update(import_finder.imports)
+
+        return nb
diff --git a/pybryt/preprocessors/intermediate_variables.py b/pybryt/preprocessors/intermediate_variables.py
index 192e5b6..3a37fb6 100644
--- a/pybryt/preprocessors/intermediate_variables.py
+++ b/pybryt/preprocessors/intermediate_variables.py
@@ -6,7 +6,9 @@ import nbformat
 
 from IPython.core.inputtransformer2 import TransformerManager
 
-from ..utils import make_secret, notebook_to_string
+from .abstract_preprocessor import AbstractPreprocessor
+
+from ..utils import make_secret
 
 
 class UnassignedVarWrapper(ast.NodeTransformer):
@@ -183,12 +185,15 @@ class UnassignedVarWrapper(ast.NodeTransformer):
         return self.transform_unassigned_node(node)
 
 
-class IntermediateVariablePreprocessor():
+class IntermediateVariablePreprocessor(AbstractPreprocessor):
     """
     Preprocessor for inserting intermediate variables in a notebook to assist in tracing.
     """
 
-    def preprocess(self, nb):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def preprocess(self, nb: nbformat.NotebookNode) -> nbformat.NotebookNode:
         """
         Preprocesses a notebook by inserting intermediate variables.
 
@@ -203,11 +208,10 @@ class IntermediateVariablePreprocessor():
         Returns:
             ``nbformat.NotebookNode``: the updated notebook
         """
-        transformer_mgr = TransformerManager()
         for cell in nb['cells']:
             if cell['cell_type'] == 'code':
                 code = cell['source']
-                code = transformer_mgr.transform_cell(code)
+                code = self.transformer_manager.transform_cell(code)
                 tree = ast.parse(code)
                 transformer = UnassignedVarWrapper()
                 transformer.add_parents(tree)


Test patch:
diff --git a/tests/annotations/test_import_.py b/tests/annotations/test_import_.py
new file mode 100644
index 0000000..7adfc59
--- /dev/null
+++ b/tests/annotations/test_import_.py
@@ -0,0 +1,118 @@
+"""Tests for import annotations"""
+
+import pytest
+
+from unittest import mock
+
+import pybryt
+from pybryt.execution.memory_footprint import MemoryFootprint
+
+from .utils import assert_object_attrs
+
+
+IMPORTED_MODULES = ["numpy", "pandas", "matplotlib"]
+
+
+def generate_import_mfp():
+    footprint = MemoryFootprint()
+    footprint.add_imports(*IMPORTED_MODULES)
+    return footprint
+
+
+@mock.patch.multiple(pybryt.ImportAnnotation, __abstractmethods__=frozenset())
+def test_import_init_behaviors():
+    """
+    """
+    with mock.patch("pybryt.annotations.import_.importlib") as mocked_importlib, \
+            mock.patch("pybryt.annotations.import_.sys") as mocked_sys:
+        pybryt.ImportAnnotation("foo")
+        mocked_importlib.import_module.assert_called_with("foo")
+        mocked_sys.modules.pop.assert_called_with("foo")
+
+        mocked_sys.modules = {"foo": None}
+        pybryt.ImportAnnotation("foo")
+        assert mocked_sys.modules == {"foo": None}
+
+        with pytest.raises(TypeError):
+            pybryt.ImportAnnotation(1)
+
+        mocked_importlib.import_module.side_effect = ImportError()
+        with pytest.raises(ValueError):
+            pybryt.ImportAnnotation("foo")
+
+
+def test_require_import():
+    """
+    """
+    footprint = generate_import_mfp()
+    pybryt.Annotation.reset_tracked_annotations()
+
+    with mock.patch("pybryt.annotations.import_.importlib") as mocked_importlib, \
+            mock.patch("pybryt.annotations.import_.sys") as mocked_sys:
+        a = pybryt.RequireImport(IMPORTED_MODULES[0])
+        b = pybryt.RequireImport("foo")
+
+    res1 = a.check(footprint)
+
+    assert_object_attrs(a, {"children__len": 0})
+    assert_object_attrs(res1, {
+        "children": [],
+        "satisfied": True,
+        "_satisfied": True,
+        "annotation": a,
+        "timestamp": -1,
+        "value": None,
+    })
+
+    assert a.to_dict() == {
+        "name": "Annotation 1",
+        "children": [],
+        "success_message": None,
+        "failure_message": None,
+        "limit": None,
+        "group": None,
+        "type": "requireimport",
+        "module": IMPORTED_MODULES[0],
+    }
+
+    res2 = b.check(footprint)
+    assert not res2.satisfied
+
+
+
+def test_forbid_import():
+    """
+    """
+    footprint = generate_import_mfp()
+    pybryt.Annotation.reset_tracked_annotations()
+
+    with mock.patch("pybryt.annotations.import_.importlib") as mocked_importlib, \
+            mock.patch("pybryt.annotations.import_.sys") as mocked_sys:
+        a = pybryt.ForbidImport("foo")
+        b = pybryt.ForbidImport(IMPORTED_MODULES[0])
+
+    res1 = a.check(footprint)
+
+    assert_object_attrs(a, {"children__len": 0})
+    assert_object_attrs(res1, {
+        "children": [],
+        "satisfied": True,
+        "_satisfied": True,
+        "annotation": a,
+        "timestamp": -1,
+        "value": None,
+    })
+
+    assert a.to_dict() == {
+        "name": "Annotation 1",
+        "children": [],
+        "success_message": None,
+        "failure_message": None,
+        "limit": None,
+        "group": None,
+        "type": "forbidimport",
+        "module": "foo",
+    }
+
+    res2 = b.check(footprint)
+    assert not res2.satisfied
diff --git a/tests/execution/test_tracing.py b/tests/execution/test_tracing.py
index 0db9641..15bcdea 100644
--- a/tests/execution/test_tracing.py
+++ b/tests/execution/test_tracing.py
@@ -145,6 +145,17 @@ def test_trace_function():
     assert np.allclose(footprint.values[0][0], arr)
     assert footprint.values[0][1] == 1
 
+    # check that imported modules are correctly tracked
+    with mock.patch.object(footprint, "add_imports") as mocked_add_imports:
+        class Foo:
+            pass
+
+        cir(frame, "return", Foo)
+        mocked_add_imports.assert_called_with(Foo.__module__.split(".")[0])
+
+        cir(frame, "return", np)
+        mocked_add_imports.assert_called_with(np.__name__)
+
 
 def test_tracing_control():
     """
diff --git a/tests/preprocessors/test_imports.py b/tests/preprocessors/test_imports.py
new file mode 100644
index 0000000..d06a92a
--- /dev/null
+++ b/tests/preprocessors/test_imports.py
@@ -0,0 +1,27 @@
+"""Tests for the imports preprocessor"""
+
+import nbformat
+
+from textwrap import dedent
+
+from pybryt.preprocessors import ImportFindingPreprocessor
+
+
+def test_preprocessor():
+    """
+    """
+    nb = nbformat.v4.new_notebook()
+    nb.cells.append(nbformat.v4.new_code_cell(dedent("""\
+        import numpy
+        import pandas as pd
+        from random import choice
+        from itertools import chain as c
+        import statsmodels.api
+        exec("import tqdm")
+    """)))
+
+    ifp = ImportFindingPreprocessor()
+
+    nb2 = ifp.preprocess(nb)
+    assert nb2 is nb
+    assert ifp.imports == {"numpy", "pandas", "random", "itertools", "statsmodels"}