Merge pull request #5000 from MarcoGorelli/hardcoded-var

fix: custom category order was hard-coded

Author: emilykl

packages/python/plotly/plotly/express/_core.py

@@ -2147,6 +2147,8 @@ def process_dataframe_timeline(args):
 
 
 def process_dataframe_pie(args, trace_patch):
+    import numpy as np
+
     names = args.get("names")
     if names is None:
         return args, trace_patch
@@ -2159,12 +2161,12 @@ def process_dataframe_pie(args, trace_patch):
     uniques = df.get_column(names).unique(maintain_order=True).to_list()
     order = [x for x in OrderedDict.fromkeys(list(order_in) + uniques) if x in uniques]
 
-    # Sort args['data_frame'] by column 'b' according to order `order`.
+    # Sort args['data_frame'] by column `names` according to order `order`.
     token = nw.generate_temporary_column_name(8, df.columns)
     args["data_frame"] = (
         df.with_columns(
-            nw.col("b")
-            .replace_strict(order, range(len(order)), return_dtype=nw.UInt32)
+            nw.col(names)
+            .replace_strict(order, np.arange(len(order)), return_dtype=nw.UInt32)
             .alias(token)
         )
         .sort(token)

packages/python/plotly/plotly/tests/test_optional/test_px/test_px_functions.py

@@ -56,6 +56,39 @@ def test_pie_like_px():
     _compare_figures(trace, fig)
 
 
+def test_pie_custom_category_order(constructor):
+    # https://github.com/plotly/plotly.py/issues/4999
+    data = {
+        "status": ["On Route", "Pending", "Waiting Result", "Delivered"],
+        "count": [28, 10, 73, 8],
+    }
+    df = constructor(data)
+    custom_order = ["Pending", "Waiting Result", "On Route", "Delivered"]
+    result = px.pie(
+        data_frame=df,
+        values="count",
+        names="status",
+        category_orders={"status": custom_order},
+    )
+    assert list(result.to_dict()["data"][0]["labels"]) == [
+        "Pending",
+        "Waiting Result",
+        "On Route",
+        "Delivered",
+    ]
+    values_ = np.array(
+        [
+            x[0]
+            for x in sorted(
+                zip(data["count"], data["status"]),
+                key=lambda t: custom_order.index(t[1]),
+            )
+        ]
+    )
+    trace = go.Pie(values=values_, labels=custom_order)
+    _compare_figures(trace, result)
+
+
 def test_sunburst_treemap_colorscales():
     labels = ["Eve", "Cain", "Seth", "Enos", "Noam", "Abel", "Awan", "Enoch", "Azura"]
     parents = ["", "Eve", "Eve", "Seth", "Seth", "Eve", "Eve", "Awan", "Eve"]