From 794ed48e1e57a1daf7b73ae9cd2c2b3260d4f3f5 Mon Sep 17 00:00:00 2001 From: alexandrusoare Date: Tue, 12 May 2026 16:09:13 +0300 Subject: [PATCH] fix(type): improve type casting --- .../chart/tool/get_chart_preview.py | 44 ++++++++++--------- 1 file changed, 24 insertions(+), 20 deletions(-) diff --git a/superset/mcp_service/chart/tool/get_chart_preview.py b/superset/mcp_service/chart/tool/get_chart_preview.py index 64758a3a255..a7dfd4edac0 100644 --- a/superset/mcp_service/chart/tool/get_chart_preview.py +++ b/superset/mcp_service/chart/tool/get_chart_preview.py @@ -56,6 +56,7 @@ from superset.mcp_service.utils.oauth2_utils import ( OAUTH2_CONFIG_ERROR_MESSAGE, ) from superset.mcp_service.utils.url_utils import get_superset_base_url +from superset.superset_typing import Column, Metric logger = logging.getLogger(__name__) @@ -145,7 +146,7 @@ class ChartLike(Protocol): uuid: Any -def _build_query_columns(form_data: Dict[str, Any]) -> list[str]: +def _build_query_columns(form_data: Dict[str, Any]) -> list[Column]: """Build query columns list from form_data, including both x_axis and groupby. Handles chart-type-specific keys: @@ -153,42 +154,45 @@ def _build_query_columns(form_data: Dict[str, Any]) -> list[str]: - Pivot tables: ``groupbyColumns`` + ``groupbyRows`` (when ``groupby`` is absent) - Mixed timeseries: ``groupby_b`` (secondary groupby) """ - x_axis_config = form_data.get("x_axis") - groupby_columns: list[Any] = form_data.get("groupby") or [] + x_axis_config: Column | None = form_data.get("x_axis") + groupby_columns: list[Column] = form_data.get("groupby") or [] # Pivot tables store dimensions under groupbyColumns / groupbyRows if not groupby_columns: - pivot_rows: list[Any] = form_data.get("groupbyRows") or [] - pivot_cols: list[Any] = form_data.get("groupbyColumns") or [] + pivot_rows: list[Column] = form_data.get("groupbyRows") or [] + pivot_cols: list[Column] = form_data.get("groupbyColumns") or [] groupby_columns = list(pivot_rows) + list(pivot_cols) # Mixed timeseries stores secondary groupby under groupby_b - groupby_b: list[Any] = form_data.get("groupby_b") or [] + groupby_b: list[Column] = form_data.get("groupby_b") or [] for col in groupby_b: if col not in groupby_columns: groupby_columns.append(col) # Deduplicate while preserving order - seen: set[Any] = set() - columns: list[Any] = [] + seen: set[str] = set() + columns: list[Column] = [] + + def _add_unique(col: Column) -> None: + key = col if isinstance(col, str) else col.get("label", str(col)) + if key not in seen: + columns.append(col) + seen.add(key) + if x_axis_config and isinstance(x_axis_config, str): - columns.append(x_axis_config) - seen.add(x_axis_config) + _add_unique(x_axis_config) elif x_axis_config and isinstance(x_axis_config, dict): col_name = x_axis_config.get("column_name") - if col_name: - columns.append(col_name) - seen.add(col_name) + if col_name and isinstance(col_name, str): + _add_unique(col_name) for col in groupby_columns: - if col not in seen: - columns.append(col) - seen.add(col) + _add_unique(col) return columns -def _build_query_metrics(form_data: Dict[str, Any]) -> list[Any]: +def _build_query_metrics(form_data: Dict[str, Any]) -> list[Metric]: """Extract metrics from form_data, handling chart-type variations. Handles: @@ -196,14 +200,14 @@ def _build_query_metrics(form_data: Dict[str, Any]) -> list[Any]: - ``metric`` (singular) — Pie charts - ``metrics_b`` — secondary y-axis in Mixed Timeseries charts """ - metrics: list[Any] = list(form_data.get("metrics") or []) + metrics: list[Metric] = list(form_data.get("metrics") or []) if not metrics: - singular = form_data.get("metric") + singular: Metric | None = form_data.get("metric") if singular: metrics = [singular] # Mixed timeseries stores the second y-axis metrics under metrics_b - metrics_b: list[Any] = form_data.get("metrics_b") or [] + metrics_b: list[Metric] = form_data.get("metrics_b") or [] for m in metrics_b: if m not in metrics: metrics.append(m)