Add outliers, label and fix rounding issue

Author: jakubno

python/e2b_code_interpreter/charts.py

@@ -165,6 +165,7 @@ class BoxAndWhiskerData:
     median: float
     third_quartile: float
     max: float
+    outliers: List[float]
 
     def __init__(self, **kwargs):
         self.label = kwargs["label"]
@@ -173,6 +174,7 @@ def __init__(self, **kwargs):
         self.median = kwargs["median"]
         self.third_quartile = kwargs["third_quartile"]
         self.max = kwargs["max"]
+        self.outliers = kwargs.get("outliers") or []
 
 
 class BoxAndWhiskerChart(Chart2D):

python/tests/graphs/test_box_and_whiskers.py

@@ -15,9 +15,6 @@
 # Create figure and axis
 fig, ax = plt.subplots(figsize=(10, 6))
 
-# Plot box plot
-ax.boxplot(data.values(), labels=data.keys())
-
 # Customize plot
 ax.set_title('Exam Scores Distribution')
 ax.set_xlabel('Class')
@@ -54,9 +51,10 @@ async def test_box_and_whiskers(async_sandbox: AsyncSandbox):
     bars = chart.elements
     assert len(bars) == 3
 
-    assert all(isinstance(bar.min, float) for bar in bars)
-    assert all(isinstance(bar.first_quartile, float) for bar in bars)
-    assert all(isinstance(bar.median, float) for bar in bars)
-    assert all(isinstance(bar.third_quartile, float) for bar in bars)
-    assert all(isinstance(bar.max, float) for bar in bars)
-    assert all(isinstance(bar.label, str) for bar in bars)
+    assert [bar.outliers for bar in bars] == [[], [76], []]
+    assert [bar.min for bar in bars] == [78, 84, 75]
+    assert [bar.first_quartile for bar in bars] == [85, 84.75, 79]
+    assert [bar.median for bar in bars] == [88, 88, 82]
+    assert [bar.third_quartile for bar in bars] == [90, 90.5, 86]
+    assert [bar.max for bar in bars] == [92, 95, 88]
+    assert [bar.label for bar in bars] == ["Class A", "Class B", "Class C"]

template/startup_scripts/0002_data.py

@@ -1,6 +1,7 @@
 from datetime import date
 import enum
 import re
+from decimal import Decimal, localcontext
 from typing import Optional, List, Tuple, Literal, Any, Union, Sequence
 
 import matplotlib
@@ -35,6 +36,18 @@ def _is_grid_line(line: Line2D) -> bool:
     return False
 
 
+def _dynamic_round(number):
+    # Convert to Decimal for precise control
+    decimal_number = Decimal(str(number))
+
+    # Dynamically determine precision based on magnitude
+    precision = max(1, 16 - decimal_number.adjusted())  # 16 digits of precision
+
+    with localcontext() as ctx:
+        ctx.prec = precision  # Set the dynamic precision
+        return +decimal_number  # The + operator applies rounding
+
+
 class ChartType(str, enum.Enum):
     LINE = "line"
     SCATTER = "scatter"
@@ -291,6 +304,7 @@ class BoxAndWhiskerData(BaseModel):
     median: float
     third_quartile: float
     max: float
+    outliers: List[float] = Field(default_factory=list)
 
 
 class BoxAndWhiskerChart(Chart2D):
@@ -301,20 +315,23 @@ class BoxAndWhiskerChart(Chart2D):
     def _extract_info(self, ax: Axes) -> None:
         super()._extract_info(ax)
 
+        labels = [item.get_text() for item in ax.get_xticklabels()]
+
         boxes = []
-        for box in ax.patches:
+        for label, box in zip(labels, ax.patches):
             vertices = box.get_path().vertices
-            x_vertices = vertices[:, 0]
-            y_vertices = vertices[:, 1]
+            x_vertices = [_dynamic_round(x) for x in vertices[:, 0]]
+            y_vertices = [_dynamic_round(y) for y in vertices[:, 1]]
             x = min(x_vertices)
             y = min(y_vertices)
             boxes.append(
                 {
                     "x": x,
                     "y": y,
-                    "label": box.get_label(),
-                    "width": round(max(x_vertices) - x, 4),
-                    "height": round(max(y_vertices) - y, 4),
+                    "label": label,
+                    "width": max(x_vertices) - x,
+                    "height": max(y_vertices) - y,
+                    "outliers": [],
                 }
             )
 
@@ -328,13 +345,21 @@ def _extract_info(self, ax: Axes) -> None:
                 box["x"], box["y"] = box["y"], box["x"]
                 box["width"], box["height"] = box["height"], box["width"]
 
-        for line in ax.lines:
-            xdata = line.get_xdata()
-            ydata = line.get_ydata()
+        for i, line in enumerate(ax.lines):
+            xdata = [_dynamic_round(x) for x in line.get_xdata()]
+            ydata = [_dynamic_round(y) for y in line.get_ydata()]
 
             if orientation == "vertical":
                 xdata, ydata = ydata, xdata
 
+            if len(xdata) == 1:
+                for box in boxes:
+                    if box["x"] <= xdata[0] <= box["x"] + box["width"]:
+                        break
+                else:
+                    continue
+
+                box["outliers"].append(ydata[0])
             if len(ydata) != 2:
                 continue
             for box in boxes:
@@ -344,6 +369,7 @@ def _extract_info(self, ax: Axes) -> None:
                 continue
 
             if (
+                # Check if the line is inside the box, prevent floating point errors
                 ydata[0] == ydata[1]
                 and box["y"] <= ydata[0] <= box["y"] + box["height"]
             ):
@@ -365,6 +391,7 @@ def _extract_info(self, ax: Axes) -> None:
                 median=box["median"],
                 third_quartile=box["y"] + box["height"],
                 max=box["whisker_upper"],
+                outliers=box["outliers"],
             )
             for box in boxes
         ]