22# Permissions are hereby granted under the terms of the MIT License:
33# https://opensource.org/licenses/MIT.
44
5- import functools
65import io
76import json
8- import operator
97import os
108import shutil
119import sys
@@ -55,6 +53,10 @@ def convert_notebook_to_script(
5553 exporter = nbconvert .PythonExporter ()
5654 (body , resources ) = exporter .from_notebook_node (self .notebook )
5755 with open (output_dir / "user_code.py" , "w" ) as fh :
56+ fh .write (
57+ "import unittest.mock\n "
58+ "get_ipython = unittest.mock.MagicMock\n "
59+ )
5860 fh .write (body )
5961 parent_dir = pathlib .Path (__file__ ).parent
6062 shutil .copy2 (parent_dir / "wrapper.py" , output_dir / "execute.py" )
@@ -73,17 +75,23 @@ def process_params_cell(self) -> None:
7375 params_cell_index = i
7476 break
7577 if params_cell_index is not None :
76- setup_code = "\n " .join (
77- map (
78- operator .attrgetter ("source" ),
79- filter (
80- lambda c : c .cell_type == "code" ,
81- self .notebook .cells [:params_cell_index ],
82- ),
83- )
78+ setup_node = nbformat .from_dict (self .notebook )
79+ setup_node .cells = setup_node .cells [:params_cell_index ]
80+ exporter = nbconvert .PythonExporter ()
81+ (setup_code , _ ) = exporter .from_notebook_node (setup_node )
82+ # Mock out the get_ipython function in case there are any
83+ # IPython magic commands in the notebook. This effectively
84+ # turns them into no-ops
85+ setup_code = (
86+ "import unittest.mock\n "
87+ "get_ipython = unittest.mock.MagicMock\n "
88+ + setup_code
8489 )
90+ params_node = nbformat .from_dict (self .notebook )
91+ params_node .cells = [params_node .cells [params_cell_index ]]
92+ (params_code , _ ) = exporter .from_notebook_node (params_node )
8593 self .nb_params = NotebookParameters .from_code (
86- self . notebook . cells [ params_cell_index ]. source ,
94+ params_code ,
8795 setup_code = setup_code ,
8896 )
8997 self .notebook .cells .insert (
0 commit comments