perf(startup): defer heavy imports in 0002_data.py

The IPython startup script imports pandas, matplotlib.pyplot, e2b_charts,
and orjson at module load time. Every Python kernel pays that cost in
resident memory, regardless of whether the customer ever displays a
DataFrame, Figure, or large JSON object — and that resident heap is
captured in the snapshot at every pause.

Move the heavy imports into the formatter __call__ bodies so they are
only resolved when a matching object is actually rendered. Top-level
imports are limited to IPython + traitlets, which is what the formatter
class definitions actually need.

IPython invokes every registered formatter for every displayed object,
so a naive "import pandas" at the top of __call__ would shift the cost
from kernel boot to the first cell with any output. Instead, gate on
sys.modules.get("pandas"): a pandas.DataFrame cannot exist unless the
user has already imported pandas. The same sentinel works for the
matplotlib + e2b_charts dependency chain in E2BChartFormatter. Sessions
that never touch pandas/matplotlib now never load them.

E2BDataFormatter previously declared type_printers={pandas.DataFrame:
…}, which forced pandas at class-definition time. The new version
dispatches via isinstance() inside __call__ — slightly more work per
formatted object, but the heavy dependencies stay unloaded until the
user genuinely needs them.

Signed-off-by: Nikita Kalyazin <nikita.kalyazin@e2b.dev>

Author: kalyazin

template/startup_scripts/0002_data.py

@@ -1,41 +1,33 @@
-import pandas
-from matplotlib.pyplot import Figure
+import sys
+
 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_}
+    def __call__(self, obj):
+        # IPython invokes every registered formatter for every displayed
+        # object. Gate on sys.modules so a non-DataFrame output (e.g. an
+        # int from `1 + 1`) doesn't pay the pandas import cost — a
+        # pandas.DataFrame can only exist if the user already imported
+        # pandas.
+        pandas = sys.modules.get("pandas")
+        if pandas is None or not isinstance(obj, pandas.DataFrame):
+            return super().__call__(obj)
+
+        result = obj.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 E2BChartFormatter(BaseFormatter):
@@ -45,19 +37,30 @@ class E2BChartFormatter(BaseFormatter):
     _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
+        # Same sys.modules gate as E2BDataFormatter: a matplotlib Figure
+        # can only exist if the user already imported matplotlib.
+        if sys.modules.get("matplotlib") is None:
+            return super().__call__(obj)
+
         from matplotlib.pyplot import Figure
 
-        if isinstance(obj, Figure):
-            return _figure_repr_e2b_chart_(obj)
-        return super().__call__(obj)
+        if not isinstance(obj, Figure):
+            return super().__call__(obj)
+
+        from e2b_charts import chart_figure_to_dict
+
+        try:
+            return chart_figure_to_dict(obj)
+        except:  # noqa: E722
+            return {}
 
 
 class E2BJSONFormatter(JSONFormatter):
     def __call__(self, obj):
         if isinstance(obj, (list, dict)):
             try:
+                import orjson
+
                 return orjson.loads(
                     orjson.dumps(
                         obj, option=orjson.OPT_SERIALIZE_NUMPY | orjson.OPT_NON_STR_KEYS