Pass imports to the manim script, fixes #11

Example.ipynb

@@ -29,6 +29,15 @@
     "from manimlib.animation.creation import ShowCreation"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import statistics"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -37,6 +46,10 @@
    "source": [
     "%%manim Shapes --low_quality\n",
     "\n",
+    "# only to demonstrate that you can use modules imported earlier\n",
+    "# (as well as variables defined beforehand, see COLOR)\n",
+    "statistics.mean([1, 2, 3])\n",
+    "\n",
     "class Shapes(Scene):\n",
     "\n",
     "    def construct(self):\n",

jupyter_manim/__init__.py

@@ -6,6 +6,7 @@
 from io import StringIO
 from pathlib import Path
 from tempfile import NamedTemporaryFile
+from types import ModuleType
 from typing import List, Dict, Tuple
 from unittest.mock import patch
 from warnings import warn
@@ -143,6 +144,25 @@ def export_globals(self):
                 f.close()
             os.remove(f.name)
 
+    def extract_imports(self):
+        frame = find_ipython_frame(inspect.stack())
+        if not frame:
+            raise Exception('Could not find IPython frame')
+
+        globals_dict = frame[0].f_globals
+
+        modules = {
+            (
+                f'import {obj.__name__} as {name}'
+                if name != obj.__name__ else
+                f'import {name}'
+            )
+            for name, obj in globals_dict.items()
+            if (not name.startswith('_')) and isinstance(obj, ModuleType)
+            if name.isidentifier() and name != 'jupyter_manim'
+        }
+        return '\n'.join(modules)
+
     def parse_arguments(self, line) -> Tuple[Dict, List]:
 
         user_args = line.split(' ')
@@ -221,6 +241,11 @@ def catch_path_and_forward(lines):
                         unpickle_script = UNPICKLE_SCRIPT.format(pickle_path=pickle_path)
                         cell = unpickle_script + cell
 
+                    try:
+                        cell = self.extract_imports() + '\n' + cell
+                    except Exception as e:
+                        warn('Assembling imports failed: ' + str(e))
+
                 f.write(cell)
                 f.close()
 

setup.py

@@ -19,7 +19,7 @@ def get_long_description(file_name):
     setup(
         name='jupyter_manim',
         packages=find_packages(),
-        version='1.2',
+        version='1.3',
         license='MIT',
         description='Cell magic rendering displaying videos in Jupyter/IPython',
         long_description=get_long_description('README.md'),

tests/test_manim.py

@@ -96,6 +96,18 @@ def mock_frame(stack):
         assert '_PICKLE_B' not in unpickled
 
 
+def test_imports(monkeypatch):
+    magics_manager = ManimMagics()
+
+    def mock_frame(stack):
+        """Return the frame of test_pickle"""
+        return inspect.stack()[0]
+
+    monkeypatch.setattr(jupyter_manim, 'find_ipython_frame', mock_frame)
+    imports = magics_manager.extract_imports()
+    assert 'import pickle' in imports
+
+
 def test_arguments_resolution():
     magics_manager = ManimMagics()
     settings, user_args = magics_manager.parse_arguments('-r 100,200')