feat: added pydantic models to store charts data (pie, bar, stacked bar, line, scatter) (#52)

* added pydantic models to store charts data (pie, bar, line)

Signed-off-by: Matteo-Omenetti <[email protected]>

* changed chart models hierarchy structure, added StackedBarChart class

Signed-off-by: Matteo-Omenetti <[email protected]>

* added scatter chart and addresed Peter's comments

Signed-off-by: Matteo-Omenetti <[email protected]>

* fixed names of classes

Signed-off-by: Matteo-Omenetti <[email protected]>

---------

Signed-off-by: Matteo-Omenetti <[email protected]>
Co-authored-by: Matteo-Omenetti <[email protected]>

Author: Matteo-Omenetti

docling_core/types/doc/document.py

@@ -100,12 +100,166 @@ class PictureMiscData(BaseModel):
     content: Dict[str, Any]
 
 
+class ChartLine(BaseModel):
+    """Represents a line in a line chart.
+
+    Attributes:
+        label (str): The label for the line.
+        values (List[Tuple[float, float]]): A list of (x, y) coordinate pairs
+            representing the line's data points.
+    """
+
+    label: str
+    values: List[Tuple[float, float]]
+
+
+class ChartBar(BaseModel):
+    """Represents a bar in a bar chart.
+
+    Attributes:
+        label (str): The label for the bar.
+        values (float): The value associated with the bar.
+    """
+
+    label: str
+    values: float
+
+
+class ChartStackedBar(BaseModel):
+    """Represents a stacked bar in a stacked bar chart.
+
+    Attributes:
+        label (List[str]): The labels for the stacked bars. Multiple values are stored
+            in cases where the chart is "double stacked," meaning bars are stacked both
+            horizontally and vertically.
+        values (List[Tuple[str, int]]): A list of values representing different segments
+            of the stacked bar along with their label.
+    """
+
+    label: List[str]
+    values: List[Tuple[str, int]]
+
+
+class ChartSlice(BaseModel):
+    """Represents a slice in a pie chart.
+
+    Attributes:
+        label (str): The label for the slice.
+        value (float): The value represented by the slice.
+    """
+
+    label: str
+    value: float
+
+
+class ChartPoint(BaseModel):
+    """Represents a point in a scatter chart.
+
+    Attributes:
+        value (Tuple[float, float]): A (x, y) coordinate pair representing a point in a
+            chart.
+    """
+
+    value: Tuple[float, float]
+
+
+class PictureChartData(BaseModel):
+    """Base class for picture chart data.
+
+    Attributes:
+        title (str): The title of the chart.
+    """
+
+    title: str
+
+
+class PictureLineChartData(PictureChartData):
+    """Represents data of a line chart.
+
+    Attributes:
+        kind (Literal["line_chart_data"]): The type of the chart.
+        x_axis_label (str): The label for the x-axis.
+        y_axis_label (str): The label for the y-axis.
+        lines (List[ChartLine]): A list of lines in the chart.
+    """
+
+    kind: Literal["line_chart_data"] = "line_chart_data"
+    x_axis_label: str
+    y_axis_label: str
+    lines: List[ChartLine]
+
+
+class PictureBarChartData(PictureChartData):
+    """Represents data of a bar chart.
+
+    Attributes:
+        kind (Literal["bar_chart_data"]): The type of the chart.
+        x_axis_label (str): The label for the x-axis.
+        y_axis_label (str): The label for the y-axis.
+        bars (List[ChartBar]): A list of bars in the chart.
+    """
+
+    kind: Literal["bar_chart_data"] = "bar_chart_data"
+    x_axis_label: str
+    y_axis_label: str
+    bars: List[ChartBar]
+
+
+class PictureStackedBarChartData(PictureChartData):
+    """Represents data of a stacked bar chart.
+
+    Attributes:
+        kind (Literal["stacked_bar_chart_data"]): The type of the chart.
+        x_axis_label (str): The label for the x-axis.
+        y_axis_label (str): The label for the y-axis.
+        stacked_bars (List[ChartStackedBar]): A list of stacked bars in the chart.
+    """
+
+    kind: Literal["stacked_bar_chart_data"] = "stacked_bar_chart_data"
+    x_axis_label: str
+    y_axis_label: str
+    stacked_bars: List[ChartStackedBar]
+
+
+class PicturePieChartData(PictureChartData):
+    """Represents data of a pie chart.
+
+    Attributes:
+        kind (Literal["pie_chart_data"]): The type of the chart.
+        slices (List[ChartSlice]): A list of slices in the pie chart.
+    """
+
+    kind: Literal["pie_chart_data"] = "pie_chart_data"
+    slices: List[ChartSlice]
+
+
+class PictureScatterChartData(PictureChartData):
+    """Represents data of a scatter chart.
+
+    Attributes:
+        kind (Literal["scatter_chart_data"]): The type of the chart.
+        x_axis_label (str): The label for the x-axis.
+        y_axis_label (str): The label for the y-axis.
+        points (List[ChartPoint]): A list of points in the scatter chart.
+    """
+
+    kind: Literal["scatter_chart_data"] = "scatter_chart_data"
+    x_axis_label: str
+    y_axis_label: str
+    points: List[ChartPoint]
+
+
 PictureDataType = Annotated[
     Union[
         PictureClassificationData,
         PictureDescriptionData,
         PictureMoleculeData,
         PictureMiscData,
+        PictureLineChartData,
+        PictureBarChartData,
+        PictureStackedBarChartData,
+        PicturePieChartData,
+        PictureScatterChartData,
     ],
     Field(discriminator="kind"),
 ]