import pandas
from matplotlib.pyplot import Figure
import IPython
from IPython.core.formatters import BaseFormatter, JSONFormatter
from traitlets.traitlets import Unicode, ObjectName
from e2b_charts import chart_figure_to_dict
import orjson
def _figure_repr_e2b_chart_(self: Figure):
"""
This method is used to extract data from the figure object to a dictionary
"""
# Get all Axes objects from the Figure
try:
return chart_figure_to_dict(self)
except: # noqa: E722
return {}
def _dataframe_repr_e2b_data_(self: pandas.DataFrame):
result = self.to_dict(orient="list")
for key, value in result.items():
# Check each column's values
result[key] = [
v.isoformat() if isinstance(v, pandas.Timestamp) else v for v in value
]
return result
class E2BDataFormatter(BaseFormatter):
format_type = Unicode("e2b/data")
print_method = ObjectName("_repr_e2b_data_")
_return_type = (dict, str)
type_printers = {pandas.DataFrame: _dataframe_repr_e2b_data_}
class E2BChartFormatter(BaseFormatter):
format_type = Unicode("e2b/chart")
print_method = ObjectName("_repr_e2b_chart_")
_return_type = (dict, str)
def __call__(self, obj):
# Figure object is for some reason removed on execution of the cell,
# so it can't be used in type_printers or with top-level import
from matplotlib.pyplot import Figure
if isinstance(obj, Figure):
return _figure_repr_e2b_chart_(obj)
return super().__call__(obj)
class E2BJSONFormatter(JSONFormatter):
def __call__(self, obj):
if isinstance(obj, (list, dict)):
try:
return orjson.loads(
orjson.dumps(
obj, option=orjson.OPT_SERIALIZE_NUMPY | orjson.OPT_NON_STR_KEYS
)
), {"expanded": True}
except TypeError:
pass
return super().__call__(obj)
ip = IPython.get_ipython()
ip.display_formatter.formatters["e2b/data"] = E2BDataFormatter(
parent=ip.display_formatter
)
ip.display_formatter.formatters["e2b/chart"] = E2BChartFormatter(
parent=ip.display_formatter
)
ip.display_formatter.formatters["application/json"] = E2BJSONFormatter(
parent=ip.display_formatter
)
