e2b-charts 0.0.1__tar.gz

e2b_charts-0.0.1/PKG-INFO

@@ -0,0 +1,25 @@
+Metadata-Version: 2.1
+Name: e2b-charts
+Version: 0.0.1
+Summary: Package for extracting data for E2B Code Interpreter
+Home-page: https://e2b.dev/
+License: Apache-2.0
+Author: e2b
+Author-email: hello@e2b.dev
+Requires-Python: >=3.10,<4.0
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: matplotlib (>=3.9.2,<4.0.0)
+Requires-Dist: numpy (>=1.26.4,<2.0.0)
+Requires-Dist: pydantic (>=2.8.2,<3.0.0)
+Project-URL: Bug Tracker, https://github.com/e2b-dev/code-interpreter/issues
+Project-URL: Repository, https://github.com/e2b-dev/e2b-code-interpreter/tree/python
+Description-Content-Type: text/markdown
+
+# Extracting Data for Code Interpreter SDK
+
+This package is a utility used to extract data in the Code Interpreter SDK from, e.g., DataFrames and matplotlib plots.
+

e2b_charts-0.0.1/README.md

@@ -0,0 +1,3 @@
+# Extracting Data for Code Interpreter SDK
+
+This package is a utility used to extract data in the Code Interpreter SDK from, e.g., DataFrames and matplotlib plots.

e2b_charts-0.0.1/e2b_charts/__init__.py

@@ -0,0 +1 @@
+from .main import graph_figure_to_graph

e2b_charts-0.0.1/e2b_charts/graphs/__init__.py

@@ -0,0 +1,4 @@
+from .base import GraphType, Graph
+from .bars import BarGraph, BoxAndWhiskerGraph
+from .pie import PieGraph
+from .planar import ScatterGraph, LineGraph

e2b_charts-0.0.1/e2b_charts/graphs/bars.py

@@ -0,0 +1,141 @@
+from typing import Literal, List
+
+from matplotlib.axes import Axes
+from pydantic import BaseModel, Field
+
+from .base import Graph2D, GraphType
+from ..utils.rounding import dynamic_round
+
+
+class BarData(BaseModel):
+    label: str
+    group: str
+    value: float
+
+
+class BarGraph(Graph2D):
+    type: Literal[GraphType.BAR] = GraphType.BAR
+
+    elements: List[BarData] = Field(default_factory=list)
+
+    def _extract_info(self, ax: Axes) -> None:
+        super()._extract_info(ax)
+        for container in ax.containers:
+            group_label = container.get_label()
+            if group_label.startswith("_container"):
+                number = int(group_label[10:])
+                group_label = f"Group {number}"
+
+            heights = [rect.get_height() for rect in container]
+            if all(height == heights[0] for height in heights):
+                # vertical bars
+                self._change_orientation()
+                labels = [label.get_text() for label in ax.get_yticklabels()]
+                values = [rect.get_width() for rect in container]
+            else:
+                # horizontal bars
+                labels = [label.get_text() for label in ax.get_xticklabels()]
+                values = heights
+            for label, value in zip(labels, values):
+
+                bar = BarData(label=label, value=value, group=group_label)
+                self.elements.append(bar)
+
+
+class BoxAndWhiskerData(BaseModel):
+    label: str
+    min: float
+    first_quartile: float
+    median: float
+    third_quartile: float
+    max: float
+    outliers: List[float]
+
+
+class BoxAndWhiskerGraph(Graph2D):
+    type: Literal[GraphType.BOX_AND_WHISKER] = GraphType.BOX_AND_WHISKER
+
+    elements: List[BoxAndWhiskerData] = Field(default_factory=list)
+
+    def _extract_info(self, ax: Axes) -> None:
+        super()._extract_info(ax)
+        labels = [item.get_text() for item in ax.get_xticklabels()]
+
+        boxes = []
+        for label, box in zip(labels, ax.patches):
+            vertices = box.get_path().vertices
+            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": label,
+                    "width": max(x_vertices) - x,
+                    "height": max(y_vertices) - y,
+                    "outliers": [],
+                }
+            )
+
+        orientation = "horizontal"
+        if all(box["height"] == boxes[0]["height"] for box in boxes):
+            orientation = "vertical"
+
+        if orientation == "vertical":
+            self._change_orientation()
+            for box in boxes:
+                box["x"], box["y"] = box["y"], box["x"]
+                box["width"], box["height"] = box["height"], box["width"]
+
+        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:
+                if box["x"] <= xdata[0] <= xdata[1] <= box["x"] + box["width"]:
+                    break
+            else:
+                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"]
+            ):
+                box["median"] = ydata[0]
+                continue
+
+            lower_value = min(ydata)
+            upper_value = max(ydata)
+            if upper_value == box["y"]:
+                box["whisker_lower"] = lower_value
+            elif lower_value == box["y"] + box["height"]:
+                box["whisker_upper"] = upper_value
+
+        self.elements = [
+            BoxAndWhiskerData(
+                label=box["label"],
+                min=box["whisker_lower"],
+                first_quartile=box["y"],
+                median=box["median"],
+                third_quartile=box["y"] + box["height"],
+                max=box["whisker_upper"],
+                outliers=box["outliers"],
+            )
+            for box in boxes
+        ]

e2b_charts-0.0.1/e2b_charts/graphs/base.py

@@ -0,0 +1,75 @@
+import enum
+import re
+from typing import Optional, List, Any
+
+from matplotlib.axes import Axes
+from pydantic import BaseModel, Field
+
+
+class GraphType(str, enum.Enum):
+    LINE = "line"
+    SCATTER = "scatter"
+    BAR = "bar"
+    PIE = "pie"
+    BOX_AND_WHISKER = "box_and_whisker"
+    SUPERGRAPH = "supergraph"
+    UNKNOWN = "unknown"
+
+
+class Graph(BaseModel):
+    type: GraphType
+    title: Optional[str] = None
+
+    elements: List[Any] = Field(default_factory=list)
+
+    def __init__(self, ax: Optional[Axes] = None, **kwargs):
+        super().__init__(**kwargs)
+        if ax:
+            self._extract_info(ax)
+
+    def _extract_info(self, ax: Axes) -> None:
+        """
+        Function to extract information for Graph
+        """
+        title = ax.get_title()
+        if title == "":
+            title = None
+
+        self.title = title
+
+
+class Graph2D(Graph):
+    x_label: Optional[str] = None
+    y_label: Optional[str] = None
+    x_unit: Optional[str] = None
+    y_unit: Optional[str] = None
+
+    def _extract_info(self, ax: Axes) -> None:
+        """
+        Function to extract information for Graph2D
+        """
+        super()._extract_info(ax)
+        x_label = ax.get_xlabel()
+        if x_label == "":
+            x_label = None
+        self.x_label = x_label
+
+        y_label = ax.get_ylabel()
+        if y_label == "":
+            y_label = None
+        self.y_label = y_label
+
+        regex = r"\s\((.*?)\)|\[(.*?)\]"
+        if self.x_label:
+            match = re.search(regex, self.x_label)
+            if match:
+                self.x_unit = match.group(1) or match.group(2)
+
+        if self.y_label:
+            match = re.search(regex, self.y_label)
+            if match:
+                self.y_unit = match.group(1) or match.group(2)
+
+    def _change_orientation(self):
+        self.x_label, self.y_label = self.y_label, self.x_label
+        self.x_unit, self.y_unit = self.y_unit, self.x_unit

e2b_charts-0.0.1/e2b_charts/graphs/pie.py

@@ -0,0 +1,32 @@
+from decimal import Decimal
+from typing import Literal, List
+
+from matplotlib.axes import Axes
+from pydantic import BaseModel, Field
+
+from .base import Graph, GraphType
+from ..utils.rounding import dynamic_round
+
+
+class PieData(BaseModel):
+    label: str
+    angle: float
+    radius: float
+
+
+class PieGraph(Graph):
+    type: Literal[GraphType.PIE] = GraphType.PIE
+
+    elements: List[PieData] = Field(default_factory=list)
+
+    def _extract_info(self, ax: Axes) -> None:
+        super()._extract_info(ax)
+
+        for wedge in ax.patches:
+            pie_data = PieData(
+                label=wedge.get_label(),
+                angle=abs(dynamic_round(Decimal(wedge.theta2) - Decimal(wedge.theta1))),
+                radius=wedge.r,
+            )
+
+            self.elements.append(pie_data)

e2b_charts-0.0.1/e2b_charts/graphs/planar.py

@@ -0,0 +1,132 @@
+from datetime import date
+from typing import List, Tuple, Union, Sequence, Any, Literal
+
+import matplotlib
+import numpy
+from matplotlib.axes import Axes
+from matplotlib.dates import _SwitchableDateConverter
+from pydantic import BaseModel, field_validator, Field
+
+from .base import Graph2D, GraphType
+from ..utils import is_grid_line
+
+
+class PointData(BaseModel):
+    label: str
+    points: List[Tuple[Union[str, float], Union[str, float]]]
+
+    @field_validator("points", mode="before")
+    @classmethod
+    def transform_points(
+        cls, value
+    ) -> List[Tuple[Union[str, float], Union[str, float]]]:
+        parsed_value = []
+        for x, y in value:
+            x = cls._parse_point(x)
+            y = cls._parse_point(y)
+            parsed_value.append((x, y))
+        return parsed_value
+
+    @staticmethod
+    def _parse_point(point):
+        if isinstance(point, date):
+            return point.isoformat()
+        if isinstance(point, numpy.datetime64):
+            return point.astype("datetime64[s]").astype(str)
+        return point
+
+
+class PointGraph(Graph2D):
+    x_ticks: List[Union[str, float]] = Field(default_factory=list)
+    x_tick_labels: List[str] = Field(default_factory=list)
+    x_scale: str = Field(default="linear")
+
+    y_ticks: List[Union[str, float]] = Field(default_factory=list)
+    y_tick_labels: List[str] = Field(default_factory=list)
+    y_scale: str = Field(default="linear")
+
+    elements: List[PointData] = Field(default_factory=list)
+
+    def _extract_info(self, ax: Axes) -> None:
+        """
+        Function to extract information for PointGraph
+        """
+        super()._extract_info(ax)
+
+        self.x_tick_labels = [label.get_text() for label in ax.get_xticklabels()]
+
+        x_ticks = ax.get_xticks()
+        self.x_ticks = self._extract_ticks_info(ax.xaxis.converter, x_ticks)
+        self.x_scale = self._detect_scale(
+            ax.xaxis.converter, ax.get_xscale(), self.x_ticks, self.x_tick_labels
+        )
+
+        self.y_tick_labels = [label.get_text() for label in ax.get_yticklabels()]
+        self.y_ticks = self._extract_ticks_info(ax.yaxis.converter, ax.get_yticks())
+        self.y_scale = self._detect_scale(
+            ax.yaxis.converter, ax.get_yscale(), self.y_ticks, self.y_tick_labels
+        )
+
+    @staticmethod
+    def _detect_scale(converter, scale: str, ticks: Sequence, labels: Sequence) -> str:
+        # If the converter is a date converter, it's a datetime scale
+        if isinstance(converter, _SwitchableDateConverter):
+            return "datetime"
+
+        # If the scale is not linear, it can't be categorical
+        if scale != "linear":
+            return scale
+
+        # If all the ticks are integers and are in order from 0 to n-1
+        # and the labels aren't corresponding to the ticks, it's categorical
+        for i, tick_and_label in enumerate(zip(ticks, labels)):
+            tick, label = tick_and_label
+            if isinstance(tick, (int, float)) and tick == i and str(i) != label:
+                continue
+            # Found a tick, which wouldn't be in a categorical scale
+            return "linear"
+
+        return "categorical"
+
+    @staticmethod
+    def _extract_ticks_info(converter: Any, ticks: Sequence) -> list:
+        if isinstance(converter, _SwitchableDateConverter):
+            return [matplotlib.dates.num2date(tick).isoformat() for tick in ticks]
+        else:
+            example_tick = ticks[0]
+
+            if isinstance(example_tick, (int, float)):
+                return [float(tick) for tick in ticks]
+            else:
+                return list(ticks)
+
+
+class LineGraph(PointGraph):
+    type: Literal[GraphType.LINE] = GraphType.LINE
+
+    def _extract_info(self, ax: Axes) -> None:
+        super()._extract_info(ax)
+
+        for line in ax.get_lines():
+            if is_grid_line(line):
+                continue
+            label = line.get_label()
+            if label.startswith("_child"):
+                number = int(label[6:])
+                label = f"Line {number}"
+
+            points = [(x, y) for x, y in zip(line.get_xdata(), line.get_ydata())]
+            line_data = PointData(label=label, points=points)
+            self.elements.append(line_data)
+
+
+class ScatterGraph(PointGraph):
+    type: Literal[GraphType.SCATTER] = GraphType.SCATTER
+
+    def _extract_info(self, ax: Axes) -> None:
+        super()._extract_info(ax)
+
+        for collection in ax.collections:
+            points = [(x, y) for x, y in collection.get_offsets()]
+            scatter_data = PointData(label=collection.get_label(), points=points)
+            self.elements.append(scatter_data)

e2b_charts-0.0.1/e2b_charts/main.py

@@ -0,0 +1,111 @@
+from typing import Optional, List, Literal
+
+from matplotlib.axes import Axes
+from matplotlib.collections import PathCollection
+from matplotlib.lines import Line2D
+from matplotlib.patches import Rectangle, Wedge, PathPatch
+from matplotlib.pyplot import Figure
+
+from matplotlib.text import Text
+from pydantic import Field
+
+from .graphs import (
+    GraphType,
+    Graph,
+    LineGraph,
+    BarGraph,
+    BoxAndWhiskerGraph,
+    PieGraph,
+    ScatterGraph,
+)
+from .utils.filtering import is_grid_line
+
+
+class SuperGraph(Graph):
+    type: Literal[GraphType.SUPERGRAPH] = GraphType.SUPERGRAPH
+    elements: List[
+        LineGraph | ScatterGraph | BarGraph | PieGraph | BoxAndWhiskerGraph
+    ] = Field(default_factory=list)
+
+    def __init__(self, figure: Figure):
+        title = figure.get_suptitle()
+        super().__init__(title=title)
+
+        self.elements = [get_graph_from_ax(ax) for ax in figure.axes]
+
+
+def _get_type_of_graph(ax: Axes) -> GraphType:
+    objects = list(filter(lambda obj: not isinstance(obj, Text), ax._children))
+
+    # Check for Line plots
+    if all(isinstance(line, Line2D) for line in objects):
+        return GraphType.LINE
+
+    if all(isinstance(box_or_path, (PathPatch, Line2D)) for box_or_path in objects):
+        return GraphType.BOX_AND_WHISKER
+
+    filtered = []
+    for obj in objects:
+        if isinstance(obj, Line2D) and is_grid_line(obj):
+            continue
+        filtered.append(obj)
+
+    objects = filtered
+
+    # Check for Scatter plots
+    if all(isinstance(path, PathCollection) for path in objects):
+        return GraphType.SCATTER
+
+    # Check for Pie plots
+    if all(isinstance(artist, Wedge) for artist in objects):
+        return GraphType.PIE
+
+    # Check for Bar plots
+    if all(isinstance(rect, Rectangle) for rect in objects):
+        return GraphType.BAR
+
+    return GraphType.UNKNOWN
+
+
+def get_graph_from_ax(
+    ax: Axes,
+) -> LineGraph | ScatterGraph | BarGraph | PieGraph | BoxAndWhiskerGraph | Graph:
+    graph_type = _get_type_of_graph(ax)
+
+    if graph_type == GraphType.LINE:
+        graph = LineGraph(ax=ax)
+    elif graph_type == GraphType.SCATTER:
+        graph = ScatterGraph(ax=ax)
+    elif graph_type == GraphType.BAR:
+        graph = BarGraph(ax=ax)
+    elif graph_type == GraphType.PIE:
+        graph = PieGraph(ax=ax)
+    elif graph_type == GraphType.BOX_AND_WHISKER:
+        graph = BoxAndWhiskerGraph(ax=ax)
+    else:
+        graph = Graph(ax=ax, type=graph_type)
+
+    return graph
+
+
+def graph_figure_to_graph(figure: Figure) -> Optional[Graph]:
+    """
+    This method is used to extract data from the figure object to a dictionary
+    """
+    # Get all Axes objects from the Figure
+    axes = figure.get_axes()
+
+    if not axes:
+        return
+    elif len(axes) > 1:
+        return SuperGraph(figure=figure)
+    else:
+        ax = axes[0]
+        return get_graph_from_ax(ax)
+
+
+def graph_figure_to_dict(figure: Figure) -> dict:
+    graph = graph_figure_to_graph(figure)
+    if graph:
+        return graph.model_dump()
+    return {}

e2b_charts-0.0.1/e2b_charts/utils/__init__.py

@@ -0,0 +1 @@
+from .filtering import is_grid_line

e2b_charts-0.0.1/e2b_charts/utils/filtering.py

@@ -0,0 +1,16 @@
+from matplotlib.lines import Line2D
+
+
+def is_grid_line(line: Line2D) -> bool:
+    x_data = line.get_xdata()
+    if len(x_data) != 2:
+        return False
+
+    y_data = line.get_ydata()
+    if len(y_data) != 2:
+        return False
+
+    if x_data[0] == x_data[1] or y_data[0] == y_data[1]:
+        return True
+
+    return False

e2b_charts-0.0.1/e2b_charts/utils/rounding.py

@@ -0,0 +1,13 @@
+from decimal import Decimal, localcontext
+
+
+def dynamic_round(number):
+    # Convert to Decimal for precise control
+    decimal_number = Decimal(str(number))
+
+    # Dynamically determine precision based on magnitude
+    precision = max(1, 8 - decimal_number.adjusted())  # 8 digits of precision
+
+    with localcontext() as ctx:
+        ctx.prec = precision  # Set the dynamic precision
+        return +decimal_number  # The + operator applies rounding

e2b_charts-0.0.1/pyproject.toml

@@ -0,0 +1,29 @@
+[tool.poetry]
+name = "e2b-charts"
+version = "0.0.1"
+description = "Package for extracting data for E2B Code Interpreter"
+authors = ["e2b <hello@e2b.dev>"]
+license = "Apache-2.0"
+readme = "README.md"
+homepage = "https://e2b.dev/"
+repository = "https://github.com/e2b-dev/e2b-code-interpreter/tree/python"
+packages = [{ include = "e2b_charts" }]
+
+[tool.poetry.dependencies]
+python = "^3.10"
+
+numpy = "^1.26.4"
+matplotlib = "^3.9.2"
+pydantic = "^2.8.2"
+
+[tool.poetry.group.dev.dependencies]
+pytest = "^7.4.0"
+python-dotenv = "^1.0.0"
+pytest-dotenv = "^0.5.2"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.poetry.urls]
+"Bug Tracker" = "https://github.com/e2b-dev/code-interpreter/issues"