feat: be able to create scatter plots

app/_types.py

@@ -1,5 +1,5 @@
 from functools import cached_property
-from typing import Annotated, Any, Optional, Tuple, cast, get_type_hints
+from typing import Annotated, Any, Literal, Optional, Tuple, Union, cast, get_type_hints
 
 import elasticsearch_dsl.query
 import luqum.tree
@@ -8,13 +8,35 @@
 
 from . import config
 from .utils import str_utils
-from .validations import check_all_values_are_fields_agg, check_index_id_is_defined
+from .validations import (
+    check_all_values_are_fields_agg,
+    check_fields_are_numeric,
+    check_index_id_is_defined,
+)
 
 #: A precise expectation of what mappings looks like in json.
 #: (dict where keys are always of type `str`).
 JSONType = dict[str, Any]
 
 
+class DistributionChartType(BaseModel):
+    """Describes an entry for a distribution chart"""
+
+    chart_type: Literal["DistributionChartType"] = "DistributionChartType"
+    field: str
+
+
+class ScatterChartType(BaseModel):
+    """Describes an entry for a scatter plot"""
+
+    chart_type: Literal["ScatterChartType"] = "ScatterChartType"
+    x: str
+    y: str
+
+
+ChartType = Union[DistributionChartType, ScatterChartType]
+
+
 class FacetItem(BaseModel):
     """Describes an entry of a facet"""
 
@@ -211,10 +233,10 @@ class SearchParameters(BaseModel):
         ),
     ] = None
     charts: Annotated[
-        list[str] | None,
+        list[ChartType] | None,
         Query(
             description="""Name of vega representations to return in the response.
-            If None (default) no charts are returned."""
+            Can be distribution chart or scatter plot"""
         ),
     ] = None
     sort_params: Annotated[
@@ -323,7 +345,34 @@ def check_facets_are_valid(self):
     @model_validator(mode="after")
     def check_charts_are_valid(self):
         """Check that the graph names are valid."""
-        errors = check_all_values_are_fields_agg(self.index_id, self.charts)
+        if self.charts is None:
+            return self
+
+        errors = check_all_values_are_fields_agg(
+            self.index_id,
+            [
+                chart.field
+                for chart in self.charts
+                if chart.chart_type == "DistributionChartType"
+            ],
+        )
+
+        errors.extend(
+            check_fields_are_numeric(
+                self.index_id,
+                [
+                    chart.x
+                    for chart in self.charts
+                    if chart.chart_type == "ScatterChartType"
+                ]
+                + [
+                    chart.y
+                    for chart in self.charts
+                    if chart.chart_type == "ScatterChartType"
+                ],
+            )
+        )
+
         if errors:
             raise ValueError(errors)
         return self

app/api.py

@@ -16,7 +16,9 @@
 from app import config
 from app._types import (
     INDEX_ID_QUERY_PARAM,
+    DistributionChartType,
     GetSearchParamsTypes,
+    ScatterChartType,
     SearchParameters,
     SearchResponse,
     SuccessSearchResponse,
@@ -107,8 +109,21 @@ def search(search_parameters: Annotated[SearchParameters, Body()]):
     return app_search.search(search_parameters)
 
 
-def parse_langs(langs: str | None) -> list[str]:
-    return langs.split(",") if langs else ["en"]
+def parse_charts_get(charts_params: str):
+    """
+    Parse for get params are 'field' or 'xfield:yfield'
+    separated by ',' for Distribution and Scatter charts.
+
+    Directly the dictionnaries in POST request
+    """
+    charts = []
+    for c in charts_params.split(","):
+        if ":" in c:
+            [x, y] = c.split(":")
+            charts.append(ScatterChartType(x=x, y=y))
+        else:
+            charts.append(DistributionChartType(field=c))
+    return charts
 
 
 @app.get("/search")
@@ -127,7 +142,7 @@ def search_get(
     langs_list = langs.split(",") if langs else ["en"]
     fields_list = fields.split(",") if fields else None
     facets_list = facets.split(",") if facets else None
-    charts_list = charts.split(",") if charts else None
+    charts_list = parse_charts_get(charts) if charts else None
     # create SearchParameters object
     try:
         search_parameters = SearchParameters(
@@ -138,8 +153,8 @@ def search_get(
             fields=fields_list,
             sort_by=sort_by,
             facets=facets_list,
-            charts=charts_list,
             index_id=index_id,
+            charts=charts_list,
         )
         return app_search.search(search_parameters)
     except ValidationError as e: