Merge pull request #476 from palewire/main

Automatically handle the `group_by_column` field on the ArrowChart

Author: palewire

datawrapper/charts/arrow.py

@@ -169,6 +169,12 @@ class ArrowChart(BaseChart):
         description="The column to label by",
     )
 
+    #: The column to group arrows by
+    groups_column: str | None = Field(
+        default=None,
+        description="The column to group arrows by",
+    )
+
     #
     # Features
     #
@@ -180,13 +186,6 @@ class ArrowChart(BaseChart):
         description="Label on the first arrow that shows column names",
     )
 
-    #: Enables the group-by-column feature, works with "Group" field
-    group_by_column: bool = Field(
-        default=False,
-        alias="group-by-column",
-        description="Enables the group-by-column feature, works with 'Group' field",
-    )
-
     @field_validator("replace_flags")
     @classmethod
     def validate_replace_flags(
@@ -224,7 +223,7 @@ def serialize_model(self) -> dict:
                 "range-extent": self.range_extent,
                 "value-label-format": self.value_label_format,
                 "color-by-column": bool(self.color_category),
-                "group-by-column": self.group_by_column,
+                "group-by-column": self.groups_column is not None,
                 "replace-flags": ReplaceFlags.serialize(self.replace_flags),
                 "show-arrow-key": self.arrow_key,
             }
@@ -240,6 +239,8 @@ def serialize_model(self) -> dict:
             axes_dict["colors"] = self.color_column
         if self.label_column is not None:
             axes_dict["labels"] = self.label_column
+        if self.groups_column is not None:
+            axes_dict["groups"] = self.groups_column
 
         # Only add axes section if there are fields to include
         if axes_dict:
@@ -319,10 +320,10 @@ def deserialize_model(cls, api_response: dict[str, Any]) -> dict[str, Any]:
             init_data["color_column"] = axes["colors"]
         if "labels" in axes:
             init_data["label_column"] = axes["labels"]
+        if "groups" in axes:
+            init_data["groups_column"] = axes["groups"]
 
         # Features
-        if "group-by-column" in visualize:
-            init_data["group_by_column"] = visualize["group-by-column"]
         if "show-arrow-key" in visualize:
             init_data["arrow_key"] = visualize["show-arrow-key"]
 

tests/integration/test_arrow_chart.py

@@ -210,7 +210,7 @@ def test_serialize_with_features(self):
             data=pd.DataFrame({"x": [1, 2], "y": [10, 20], "z": [15, 25]}),
             start_column="y",
             end_column="z",
-            group_by_column=True,
+            groups_column="x",  # Setting groups_column should enable group-by-column
             arrow_key=True,
         )
 
@@ -289,7 +289,7 @@ def mock_get(url):
             # Verify features
             assert chart.thick_arrows is True
             assert chart.arrow_key is True
-            assert chart.group_by_column is True
+            assert chart.groups_column == "Gender"
 
             # Verify sorting
             assert chart.sort_ranges is False