import ast import base64 import importlib import io import os import re import traceback from typing import Any, List import matplotlib.pyplot as plt import pandas as pd import plotly.io as pio from lida_ko.datamodel import ChartExecutorResponse, Summary def preprocess_code(code: str) -> str: """Preprocess code to remove any preamble and explanation text""" code = code.replace("", "") code = code.replace("", "") code = code.replace("", "") # remove all text after chart = plot(data) if "chart = plot(data)" in code: # print(code) index = code.find("chart = plot(data)") if index != -1: code = code[: index + len("chart = plot(data)")] if "```" in code: pattern = r"```(?:\w+\n)?([\s\S]+?)```" matches = re.findall(pattern, code) if matches: code = matches[0] # code = code.replace("```", "") # return code if "import" in code: # return only text after the first import statement index = code.find("import") if index != -1: code = code[index:] code = code.replace("```", "") if "chart = plot(data)" not in code: code = code + "\nchart = plot(data)" return code def get_globals_dict(code_string, data): # Parse the code string into an AST tree = ast.parse(code_string) # Extract the names of the imported modules and their aliases imported_modules = [] for node in tree.body: if isinstance(node, ast.Import): for alias in node.names: module = importlib.import_module(alias.name) imported_modules.append((alias.name, alias.asname, module)) elif isinstance(node, ast.ImportFrom): module = importlib.import_module(node.module) for alias in node.names: obj = getattr(module, alias.name) imported_modules.append( (f"{node.module}.{alias.name}", alias.asname, obj) ) # Import the required modules into a dictionary globals_dict = {} for module_name, alias, obj in imported_modules: if alias: globals_dict[alias] = obj else: globals_dict[module_name.split(".")[-1]] = obj ex_dicts = {"pd": pd, "data": data, "plt": plt} globals_dict.update(ex_dicts) return globals_dict class ChartExecutor: """Execute code and return chart object""" def __init__(self) -> None: pass def execute( self, code_specs: List[str], data: Any, summary: Summary, library="altair", return_error: bool = False, ) -> Any: """Validate and convert code""" # # check if user has given permission to execute code. if env variable # # LIDA_ALLOW_CODE_EVAL is set to '1'. Else raise exception # if os.environ.get("LIDA_ALLOW_CODE_EVAL") != '1': # raise Exception( # "Permission to execute code not granted. Please set the environment variable LIDA_ALLOW_CODE_EVAL to '1' to allow code execution.") if isinstance(summary, dict): summary = Summary(**summary) charts = [] code_spec_copy = code_specs.copy() code_specs = [preprocess_code(code) for code in code_specs] if library == "altair": for code in code_specs: try: ex_locals = get_globals_dict(code, data) exec(code, ex_locals) chart = ex_locals["chart"] vega_spec = chart.to_dict() del vega_spec["data"] if "datasets" in vega_spec: del vega_spec["datasets"] vega_spec["data"] = {"url": f"/files/data/{summary.file_name}"} charts.append( ChartExecutorResponse( spec=vega_spec, status=True, raster=None, code=code, library=library, ) ) except Exception as exception_error: print(code_spec_copy, "\n===========\n") print(exception_error) print(traceback.format_exc()) if return_error: charts.append( ChartExecutorResponse( spec=None, status=False, raster=None, code=code, library=library, error={ "message": str(exception_error), "traceback": traceback.format_exc(), }, ) ) return charts elif library == "matplotlib" or library == "seaborn": # print colum dtypes for code in code_specs: try: ex_locals = get_globals_dict(code, data) # print(ex_locals) exec(code, ex_locals) chart = ex_locals["chart"] if plt: buf = io.BytesIO() plt.box(False) plt.grid(color="lightgray", linestyle="dashed", zorder=-10) # try: # plt.draw() # # plt.tight_layout() # except AttributeError: # print("Warning: tight_layout encountered an error. The layout may not be optimal.") # pass plt.savefig(buf, format="png", dpi=100, pad_inches=0.2) buf.seek(0) plot_data = base64.b64encode(buf.read()).decode("ascii") plt.close() charts.append( ChartExecutorResponse( spec=None, status=True, raster=plot_data, code=code, library=library, ) ) except Exception as exception_error: print(code_spec_copy[0]) print("****\n", str(exception_error)) # print(traceback.format_exc()) if return_error: charts.append( ChartExecutorResponse( spec=None, status=False, raster=None, code=code, library=library, error={ "message": str(exception_error), "traceback": traceback.format_exc(), }, ) ) return charts elif library == "ggplot": # print colum dtypes for code in code_specs: try: ex_locals = get_globals_dict(code, data) exec(code, ex_locals) chart = ex_locals["chart"] if plt: buf = io.BytesIO() chart.save(buf, format="png") plot_data = base64.b64encode(buf.getvalue()).decode("utf-8") charts.append( ChartExecutorResponse( spec=None, status=True, raster=plot_data, code=code, library=library, ) ) except Exception as exception_error: print(code) print(traceback.format_exc()) if return_error: charts.append( ChartExecutorResponse( spec=None, status=False, raster=None, code=code, library=library, error={ "message": str(exception_error), "traceback": traceback.format_exc(), }, ) ) return charts elif library == "plotly": for code in code_specs: try: ex_locals = get_globals_dict(code, data) exec(code, ex_locals) chart = ex_locals["chart"] if pio: chart_bytes = pio.to_image(chart, 'png') plot_data = base64.b64encode(chart_bytes).decode('utf-8') charts.append( ChartExecutorResponse( spec=None, status=True, raster=plot_data, code=code, library=library, ) ) except Exception as exception_error: print(code) print(traceback.format_exc()) if return_error: charts.append( ChartExecutorResponse( spec=None, status=False, raster=None, code=code, library=library, error={ "message": str(exception_error), "traceback": traceback.format_exc(), }, ) ) return charts else: raise Exception( f"Unsupported library. Supported libraries are altair, matplotlib, seaborn, ggplot, plotly. You provided {library}" )