Add units + improve graph type detection

Author: jakubno

python/tests/graphs/test_bar.py

@@ -36,6 +36,9 @@ async def test_graph_bar(async_sandbox: AsyncCodeInterpreter):
     assert graphs[0]['x_label'] == "Authors"
     assert graphs[0]['y_label'] == "Number of Books Sold"
 
+    assert graphs[0]['x_unit'] is None
+    assert graphs[0]['y_unit'] is None
+
     assert all(isinstance(x, int) for x in graphs[0]['x_ticks'])
     assert all(isinstance(y, float) for y in graphs[0]['y_ticks'])
 
@@ -53,4 +56,3 @@ async def test_graph_bar(async_sandbox: AsyncCodeInterpreter):
     assert all(isinstance(x, (int, float)) for x in data['heights'])
     assert all(isinstance(y, (int, float)) for y in data['x'])
     assert all(isinstance(y, (int, float)) for y in data['y'])
-

python/tests/graphs/test_line.py

@@ -17,8 +17,8 @@
 plt.plot(x, y_cos, label='cos(x)')
 
 # Add labels and title
-plt.xlabel('x')
-plt.ylabel('y')
+plt.xlabel("Time (s)")
+plt.ylabel("Amplitude (Hz)")
 plt.title('Plot of sin(x) and cos(x)')
 
 # Display the plot
@@ -37,8 +37,10 @@ async def test_line_graph(async_sandbox: AsyncCodeInterpreter):
 
     assert graphs[0]['type'] == "line"
     assert graphs[0]['title'] == "Plot of sin(x) and cos(x)"
-    assert graphs[0]['x_label'] == "x"
-    assert graphs[0]['y_label'] == "y"
+    assert graphs[0]['x_label'] == "Time (s)"
+    assert graphs[0]['y_label'] == "Amplitude (Hz)"
+    assert graphs[0]['x_unit'] == "s"
+    assert graphs[0]['y_unit'] == "Hz"
 
     assert all(isinstance(x, float) for x in graphs[0]['x_ticks'])
     assert all(isinstance(y, float) for y in graphs[0]['y_ticks'])

python/tests/graphs/test_pie.py

@@ -39,6 +39,8 @@ async def test_pie_graph(async_sandbox: AsyncCodeInterpreter):
     assert graphs[0]['title'] == "Will I wake up early tomorrow?"
     assert graphs[0]['x_label'] == "x"
     assert graphs[0]['y_label'] == "y"
+    assert graphs[0]['x_unit'] is None
+    assert graphs[0]['y_unit'] is None
 
     assert len(graphs[0]['x_ticks']) == 0
     assert len(graphs[0]['y_ticks']) == 0

template/startup_scripts/0002_data.py

@@ -1,14 +1,17 @@
 import enum
+import re
+from typing import Optional
 
 import pandas
 from matplotlib.axes import Axes
 from matplotlib.collections import PathCollection
 from matplotlib.lines import Line2D
-from matplotlib.patches import Rectangle, Wedge
+from matplotlib.patches import Rectangle, Wedge, PathPatch
 from matplotlib.pyplot import Figure
 import IPython
 
 from IPython.core.formatters import BaseFormatter
+from matplotlib.text import Text
 from traitlets.traitlets import Unicode, ObjectName
 
 
@@ -20,21 +23,34 @@ class PlotType(enum.Enum):
     UNKNOWN = "unknown"
 
 
-def get_type_of_plot(ax: Axes) -> PlotType:
+def _extract_units(label: str) -> Optional[str]:
+    """
+    Function to extract units from labels
+    """
+    # Look for units in parentheses or brackets
+    match = re.search(r"\s\((.*?)\)|\[(.*?)\]", label)
+    if match:
+        return match.group(1) or match.group(2)  # return the matched unit
+    return None  # No units found
+
+
+def _get_type_of_plot(ax: Axes) -> PlotType:
+    objects = list(filter(lambda obj: not isinstance(obj, Text), ax._children))
+
     # Check for Line plots
-    if any(isinstance(line, Line2D) for line in ax.get_lines()):
+    if all(isinstance(line, Line2D) for line in objects):
         return PlotType.LINE
 
     # Check for Scatter plots
-    if any(isinstance(collection, PathCollection) for collection in ax.collections):
+    if all(isinstance(path, PathCollection) for path in objects):
         return PlotType.SCATTER
 
     # Check for Pie plots
-    if any(isinstance(artist, Wedge) for artist in ax.patches):
+    if all(isinstance(artist, Wedge) for artist in objects):
         return PlotType.PIE
 
     # Check for Bar plots
-    if any(isinstance(rect, Rectangle) for rect in ax.patches):
+    if all(isinstance(rect, Rectangle) for rect in objects):
         return PlotType.BAR
 
     return PlotType.UNKNOWN
@@ -53,17 +69,19 @@ def _figure_repr_e2b_data_(self: Figure):
         ax_data = {
             "title": ax.get_title(),
             "x_label": ax.get_xlabel(),
+            "x_unit": _extract_units(ax.get_xlabel()),
             "x_ticks": ax.get_xticks(),
             "x_tick_labels": [label.get_text() for label in ax.get_xticklabels()],
             "x_scale": ax.get_xscale(),
             "y_label": ax.get_ylabel(),
+            "y_unit": _extract_units(ax.get_ylabel()),
             "y_ticks": ax.get_yticks(),
             "y_tick_labels": [label.get_text() for label in ax.get_yticklabels()],
             "y_scale": ax.get_yscale(),
             "data": [],
         }
 
-        plot_type = get_type_of_plot(ax)
+        plot_type = _get_type_of_plot(ax)
         ax_data["type"] = plot_type.value
 
         if plot_type == PlotType.LINE: