83 lines
2.3 KiB
Python
83 lines
2.3 KiB
Python
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
|
|
)
|