Skip to content

Commit

Permalink
feat: Add Heatmap chart migration logic
Browse files Browse the repository at this point in the history
  • Loading branch information
michael-s-molina committed Mar 29, 2024
1 parent c0f8dfc commit aad847b
Show file tree
Hide file tree
Showing 3 changed files with 96 additions and 0 deletions.
3 changes: 3 additions & 0 deletions superset/cli/viz_migrations.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ class VizType(str, Enum):
BUBBLE = "bubble"
DUAL_LINE = "dual_line"
LINE = "line"
HEATMAP = "heatmap"
PIVOT_TABLE = "pivot_table"
SUNBURST = "sunburst"
TREEMAP = "treemap"
Expand Down Expand Up @@ -80,6 +81,7 @@ def migrate(viz_type: VizType, is_downgrade: bool = False) -> None:
MigrateBubbleChart,
MigrateDualLine,
MigrateLineChart,
MigrateHeatmapChart,
MigratePivotTable,
MigrateSunburst,
MigrateTreeMap,
Expand All @@ -90,6 +92,7 @@ def migrate(viz_type: VizType, is_downgrade: bool = False) -> None:
VizType.BUBBLE: MigrateBubbleChart,
VizType.DUAL_LINE: MigrateDualLine,
VizType.LINE: MigrateLineChart,
VizType.HEATMAP: MigrateHeatmapChart,
VizType.PIVOT_TABLE: MigratePivotTable,
VizType.SUNBURST: MigrateSunburst,
VizType.TREEMAP: MigrateTreeMap,
Expand Down
15 changes: 15 additions & 0 deletions superset/migrations/shared/migrate_viz/processors.py
Original file line number Diff line number Diff line change
Expand Up @@ -213,3 +213,18 @@ def _pre_action(self) -> None:

# Truncate y-axis by default to preserve layout
self.data["y_axis_showminmax"] = True


class MigrateHeatmapChart(MigrateViz):
source_viz_type = "heatmap"
target_viz_type = "heatmap_v2"
rename_keys = {
"all_columns_x": "x_axis",
"all_columns_y": "groupby",
"y_axis_bounds": "value_bounds",
"show_perc": "show_percentage",
}
remove_keys = {"sort_by_metric", "canvas_image_rendering"}

def _pre_action(self) -> None:
self.data["legend_type"] = "continuous"
78 changes: 78 additions & 0 deletions tests/unit_tests/migrations/viz/heatmap_v1_v2_test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
from typing import Any

from superset.migrations.shared.migrate_viz import MigrateHeatmapChart
from tests.unit_tests.migrations.viz.utils import migrate_and_assert

SOURCE_FORM_DATA: dict[str, Any] = {
"any_other_key": "untouched",
"all_columns_x": ["category"],
"all_columns_y": ["product"],
"metric": ["sales"],
"adhoc_filters": [],
"row_limit": 100,
"sort_by_metric": True,
"linear_color_scheme": "blue",
"xscale_interval": 2,
"yscale_interval": 2,
"canvas_image_rendering": "auto",
"normalize_across": "x",
"left_margin": 50,
"bottom_margin": 50,
"y_axis_bounds": [0, 100],
"y_axis_format": "SMART_NUMBER",
"currency_format": "USD",
"sort_x_axis": "alpha_asc",
"sort_y_axis": "alpha_asc",
"show_legend": True,
"show_perc": True,
"show_values": True,
"normalized": True,
"viz_type": "heatmap",
}

TARGET_FORM_DATA: dict[str, Any] = {
"any_other_key": "untouched",
"x_axis": ["category"],
"groupby": ["product"],
"metric": ["sales"],
"adhoc_filters": [],
"row_limit": 100,
"legend_type": "continuous",
"linear_color_scheme": "blue",
"xscale_interval": 2,
"yscale_interval": 2,
"normalize_across": "x",
"left_margin": 50,
"bottom_margin": 50,
"value_bounds": [0, 100],
"y_axis_format": "SMART_NUMBER",
"currency_format": "USD",
"sort_x_axis": "alpha_asc",
"sort_y_axis": "alpha_asc",
"show_legend": True,
"show_percentage": True,
"show_values": True,
"normalized": True,
"viz_type": "heatmap_v2",
"form_data_bak": SOURCE_FORM_DATA,
}


def test_migration() -> None:
migrate_and_assert(MigrateHeatmapChart, SOURCE_FORM_DATA, TARGET_FORM_DATA)

0 comments on commit aad847b

Please sign in to comment.