Rename to charts

Author: jakubno

chart_data_extractor/e2b_charts/__init__.py

@@ -1 +1 @@
-from .main import graph_figure_to_graph
+from .main import chart_figure_to_chart, chart_figure_to_dict

chart_data_extractor/e2b_charts/charts/__init__.py

@@ -0,0 +1,4 @@
+from .base import ChartType, Chart
+from .bars import BarChart, BoxAndWhiskerChart
+from .pie import PieChart
+from .planar import ScatterChart, LineChart

chart_data_extractor/e2b_charts/graphs/bars.py renamed to chart_data_extractor/e2b_charts/charts/bars.py

@@ -3,7 +3,7 @@
 from matplotlib.axes import Axes
 from pydantic import BaseModel, Field
 
-from .base import Graph2D, GraphType
+from .base import Chart2D, ChartType
 from ..utils.rounding import dynamic_round
 
 
@@ -13,8 +13,8 @@ class BarData(BaseModel):
     value: float
 
 
-class BarGraph(Graph2D):
-    type: Literal[GraphType.BAR] = GraphType.BAR
+class BarChart(Chart2D):
+    type: Literal[ChartType.BAR] = ChartType.BAR
 
     elements: List[BarData] = Field(default_factory=list)
 
@@ -52,8 +52,8 @@ class BoxAndWhiskerData(BaseModel):
     outliers: List[float]
 
 
-class BoxAndWhiskerGraph(Graph2D):
-    type: Literal[GraphType.BOX_AND_WHISKER] = GraphType.BOX_AND_WHISKER
+class BoxAndWhiskerChart(Chart2D):
+    type: Literal[ChartType.BOX_AND_WHISKER] = ChartType.BOX_AND_WHISKER
 
     elements: List[BoxAndWhiskerData] = Field(default_factory=list)
 

chart_data_extractor/e2b_charts/graphs/base.py renamed to chart_data_extractor/e2b_charts/charts/base.py

@@ -6,18 +6,18 @@
 from pydantic import BaseModel, Field
 
 
-class GraphType(str, enum.Enum):
+class ChartType(str, enum.Enum):
     LINE = "line"
     SCATTER = "scatter"
     BAR = "bar"
     PIE = "pie"
     BOX_AND_WHISKER = "box_and_whisker"
-    SUPERGRAPH = "supergraph"
+    SUPERCHART = "superchart"
     UNKNOWN = "unknown"
 
 
-class Graph(BaseModel):
-    type: GraphType
+class Chart(BaseModel):
+    type: ChartType
     title: Optional[str] = None
 
     elements: List[Any] = Field(default_factory=list)
@@ -29,7 +29,7 @@ def __init__(self, ax: Optional[Axes] = None, **kwargs):
 
     def _extract_info(self, ax: Axes) -> None:
         """
-        Function to extract information for Graph
+        Function to extract information for Chart
         """
         title = ax.get_title()
         if title == "":
@@ -38,15 +38,15 @@ def _extract_info(self, ax: Axes) -> None:
         self.title = title
 
 
-class Graph2D(Graph):
+class Chart2D(Chart):
     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
+        Function to extract information for Chart2D
         """
         super()._extract_info(ax)
         x_label = ax.get_xlabel()

chart_data_extractor/e2b_charts/graphs/pie.py renamed to chart_data_extractor/e2b_charts/charts/pie.py

@@ -4,7 +4,7 @@
 from matplotlib.axes import Axes
 from pydantic import BaseModel, Field
 
-from .base import Graph, GraphType
+from .base import Chart, ChartType
 from ..utils.rounding import dynamic_round
 
 
@@ -14,8 +14,8 @@ class PieData(BaseModel):
     radius: float
 
 
-class PieGraph(Graph):
-    type: Literal[GraphType.PIE] = GraphType.PIE
+class PieChart(Chart):
+    type: Literal[ChartType.PIE] = ChartType.PIE
 
     elements: List[PieData] = Field(default_factory=list)
 

chart_data_extractor/e2b_charts/graphs/planar.py renamed to chart_data_extractor/e2b_charts/charts/planar.py

@@ -7,7 +7,7 @@
 from matplotlib.dates import _SwitchableDateConverter
 from pydantic import BaseModel, field_validator, Field
 
-from .base import Graph2D, GraphType
+from .base import Chart2D, ChartType
 from ..utils import is_grid_line
 
 
@@ -36,7 +36,7 @@ def _parse_point(point):
         return point
 
 
-class PointGraph(Graph2D):
+class PointChart(Chart2D):
     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")
@@ -49,7 +49,7 @@ class PointGraph(Graph2D):
 
     def _extract_info(self, ax: Axes) -> None:
         """
-        Function to extract information for PointGraph
+        Function to extract information for PointChart
         """
         super()._extract_info(ax)
 
@@ -101,8 +101,8 @@ def _extract_ticks_info(converter: Any, ticks: Sequence) -> list:
                 return list(ticks)
 
 
-class LineGraph(PointGraph):
-    type: Literal[GraphType.LINE] = GraphType.LINE
+class LineChart(PointChart):
+    type: Literal[ChartType.LINE] = ChartType.LINE
 
     def _extract_info(self, ax: Axes) -> None:
         super()._extract_info(ax)
@@ -120,8 +120,8 @@ def _extract_info(self, ax: Axes) -> None:
             self.elements.append(line_data)
 
 
-class ScatterGraph(PointGraph):
-    type: Literal[GraphType.SCATTER] = GraphType.SCATTER
+class ScatterChart(PointChart):
+    type: Literal[ChartType.SCATTER] = ChartType.SCATTER
 
     def _extract_info(self, ax: Axes) -> None:
         super()._extract_info(ax)

chart_data_extractor/e2b_charts/graphs/__init__.py

@@ -9,40 +9,40 @@
 from matplotlib.text import Text
 from pydantic import Field
 
-from .graphs import (
-    GraphType,
-    Graph,
-    LineGraph,
-    BarGraph,
-    BoxAndWhiskerGraph,
-    PieGraph,
-    ScatterGraph,
+from .charts import (
+    ChartType,
+    Chart,
+    LineChart,
+    BarChart,
+    BoxAndWhiskerChart,
+    PieChart,
+    ScatterChart,
 )
 from .utils.filtering import is_grid_line
 
 
-class SuperGraph(Graph):
-    type: Literal[GraphType.SUPERGRAPH] = GraphType.SUPERGRAPH
+class SuperChart(Chart):
+    type: Literal[ChartType.SUPERCHART] = ChartType.SUPERCHART
     elements: List[
-        LineGraph | ScatterGraph | BarGraph | PieGraph | BoxAndWhiskerGraph
+        LineChart | ScatterChart | BarChart | PieChart | BoxAndWhiskerChart
     ] = 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]
+        self.elements = [get_chart_from_ax(ax) for ax in figure.axes]
 
 
-def _get_type_of_graph(ax: Axes) -> GraphType:
+def _get_type_of_chart(ax: Axes) -> ChartType:
     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
+        return ChartType.LINE
 
     if all(isinstance(box_or_path, (PathPatch, Line2D)) for box_or_path in objects):
-        return GraphType.BOX_AND_WHISKER
+        return ChartType.BOX_AND_WHISKER
 
     filtered = []
     for obj in objects:
@@ -54,41 +54,41 @@ def _get_type_of_graph(ax: Axes) -> GraphType:
 
     # Check for Scatter plots
     if all(isinstance(path, PathCollection) for path in objects):
-        return GraphType.SCATTER
+        return ChartType.SCATTER
 
     # Check for Pie plots
     if all(isinstance(artist, Wedge) for artist in objects):
-        return GraphType.PIE
+        return ChartType.PIE
 
     # Check for Bar plots
     if all(isinstance(rect, Rectangle) for rect in objects):
-        return GraphType.BAR
+        return ChartType.BAR
 
-    return GraphType.UNKNOWN
+    return ChartType.UNKNOWN
 
 
-def get_graph_from_ax(
+def get_chart_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)
+) -> LineChart | ScatterChart | BarChart | PieChart | BoxAndWhiskerChart | Chart:
+    chart_type = _get_type_of_chart(ax)
+
+    if chart_type == ChartType.LINE:
+        chart = LineChart(ax=ax)
+    elif chart_type == ChartType.SCATTER:
+        chart = ScatterChart(ax=ax)
+    elif chart_type == ChartType.BAR:
+        chart = BarChart(ax=ax)
+    elif chart_type == ChartType.PIE:
+        chart = PieChart(ax=ax)
+    elif chart_type == ChartType.BOX_AND_WHISKER:
+        chart = BoxAndWhiskerChart(ax=ax)
     else:
-        graph = Graph(ax=ax, type=graph_type)
+        chart = Chart(ax=ax, type=chart_type)
 
-    return graph
+    return chart
 
 
-def graph_figure_to_graph(figure: Figure) -> Optional[Graph]:
+def chart_figure_to_chart(figure: Figure) -> Optional[Chart]:
     """
     This method is used to extract data from the figure object to a dictionary
     """
@@ -98,14 +98,14 @@ def graph_figure_to_graph(figure: Figure) -> Optional[Graph]:
     if not axes:
         return
     elif len(axes) > 1:
-        return SuperGraph(figure=figure)
+        return SuperChart(figure=figure)
     else:
         ax = axes[0]
-        return get_graph_from_ax(ax)
+        return get_chart_from_ax(ax)
 
 
-def graph_figure_to_dict(figure: Figure) -> dict:
-    graph = graph_figure_to_graph(figure)
-    if graph:
-        return graph.model_dump()
+def chart_figure_to_dict(figure: Figure) -> dict:
+    chart = chart_figure_to_chart(figure)
+    if chart:
+        return chart.model_dump()
     return {}

chart_data_extractor/e2b_charts/main.py

@@ -1,42 +1,42 @@
 import matplotlib.pyplot as plt
 
-from e2b_charts import graph_figure_to_graph
-from e2b_charts.graphs import BarGraph, GraphType
+from e2b_charts import chart_figure_to_chart
+from e2b_charts.charts import BarChart, ChartType
 
 
-def _prep_graph_figure():
+def _prep_chart_figure():
     # Prepare data
     authors = ["Author A", "Author B", "Author C", "Author D"]
     sales = [100, 200, 300, 400]
 
-    # Create and customize the bar graph
+    # Create and customize the bar chart
     plt.figure(figsize=(10, 6))
     plt.bar(authors, sales, label="Books Sold", color="blue")
     plt.xlabel("Authors")
     plt.ylabel("Number of Books Sold")
     plt.title("Book Sales by Authors")
 
-    # Display the graph
+    # Display the chart
     plt.tight_layout()
     return plt.gcf()
 
 
-def test_graph_bar():
-    figure = _prep_graph_figure()
-    graph = graph_figure_to_graph(figure)
-    assert graph
+def test_chart_bar():
+    figure = _prep_chart_figure()
+    chart = chart_figure_to_chart(figure)
+    assert chart
 
-    assert isinstance(graph, BarGraph)
-    assert graph.type == GraphType.BAR
-    assert graph.title == "Book Sales by Authors"
+    assert isinstance(chart, BarChart)
+    assert chart.type == ChartType.BAR
+    assert chart.title == "Book Sales by Authors"
 
-    assert graph.x_label == "Authors"
-    assert graph.y_label == "Number of Books Sold"
+    assert chart.x_label == "Authors"
+    assert chart.y_label == "Number of Books Sold"
 
-    assert graph.x_unit is None
-    assert graph.y_unit is None
+    assert chart.x_unit is None
+    assert chart.y_unit is None
 
-    bars = graph.elements
+    bars = chart.elements
     assert len(bars) == 4
 
     assert [bar.value for bar in bars] == [100, 200, 300, 400]