run_code.py

import contextlib
import io
import os
import sys
import tempfile
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from IPython.core.interactiveshell import InteractiveShell



class InteractiveShellManager:
    def __init__(self):
        self.shell = InteractiveShell.instance()

    def run_code(self, code):
        # Create a temporary directory to save the pictures generated by matplotlib
        with tempfile.TemporaryDirectory() as tempdir:
            # Set the picture save path of matplotlib
            plt.rcParams['savefig.directory'] = "./temp"

            # Create a context manager to capture stdout
            @contextlib.contextmanager
            def stdout_capture():
                old_stdout = sys.stdout
                sys.stdout = io.StringIO()
                yield sys.stdout
                sys.stdout = old_stdout

            with stdout_capture() as new_stdout:
                result = self.shell.run_cell(code)

            tempdir = "./images"

            if result.success:
                print_out = new_stdout.getvalue()
                result_out = result.result
                image_files = os.listdir(tempdir)
                image_paths = [os.path.join(tempdir, img) for img in image_files if img.endswith('.png')]

                return {
                    "print": print_out,
                    "result": result_out,
                    "images": image_paths
                }
            else:
                raise Exception('error: ' + str(result.error_in_exec))

    def close(self):
        self.shell = None