handle more cases

emmanuelle · emmanuelle · commit 4e93aa73899d · 2020-07-15T11:17:56.000+02:00
diff --git a/packages/python/plotly/plotly/express/_core.py b/packages/python/plotly/plotly/express/_core.py
@@ -1052,6 +1052,15 @@ def process_args_into_dataframe(args, wide_mode, var_name, value_name):
         else:
             df_output[df_input.columns] = df_input[df_input.columns]
 
+    # Case of sunburst and treemap, used with the `path` argument
+    # in which case the dataframe is massaged into a new format
+    # and column with reserved names will be created
+    uses_path = False
+    reserved_names = []
+    if "path" in args and args["path"] is not None:
+        uses_path = True
+        reserved_names = ["id", "labels", "parent"]
+
     # hover_data is a dict
     hover_data_is_dict = (
         "hover_data" in args
@@ -1095,6 +1104,7 @@ def process_args_into_dataframe(args, wide_mode, var_name, value_name):
             if argument is None:
                 continue
             col_name = None
+
             # Case of multiindex
             if isinstance(argument, pd.MultiIndex):
                 raise TypeError(
@@ -1114,7 +1124,19 @@ def process_args_into_dataframe(args, wide_mode, var_name, value_name):
                     ranges.append(col_name)
             # ----------------- argument is likely a col name ----------------------
             elif isinstance(argument, str) or not hasattr(argument, "__len__"):
-                if (
+                if uses_path and argument in reserved_names and field_name != 'path':
+                    if (args["labels"] is None or argument not in args["labels"]):
+                        raise ValueError(
+                            "%s is a reserved name for px.sunburst and px.treemap. "
+                            "Please use the labels argument to provide another name "
+                            "for the column, for example "
+                            "labels={'%s': '%s_col'}"
+                            % (argument, argument, argument)
+                            )
+                    else:
+                        col_name = args["labels"][argument]
+                        df_output[col_name] = df_input[argument]
+                elif (
                     field_name == "hover_data"
                     and hover_data_is_dict
                     and args["hover_data"][str(argument)][1] is not None
@@ -1484,12 +1506,15 @@ def process_dataframe_hierarchy(args):
     _check_dataframe_all_leaves(df[path[::-1]])
     discrete_color = False
 
+    reserved_names = ["labels", "parent", "id"]
+
     new_path = []
     for col_name in path:
         new_col_name = col_name + "_path_copy"
         new_path.append(new_col_name)
         df[new_col_name] = df[col_name]
     path = new_path
+
     # ------------ Define aggregation functions --------------------------------
 
     def aggfunc_discrete(x):
@@ -1547,9 +1572,14 @@ def aggfunc_continuous(x):
         if col not in agg_f:
             agg_f[col] = aggfunc_discrete
     # Avoid collisions with reserved names - columns in the path have been copied already
-    cols = list(set(cols) - set(["labels", "parent", "id"]))
+    #for col in cols:
+    #    if col in reserved_names and args["labels"] is not None and col in args["labels"]:
+    #        df[args["labels"][col]] = df[col]
+    cols = list(set(cols) - set(reserved_names))
+
+
     # ----------------------------------------------------------------------------
-    df_all_trees = pd.DataFrame(columns=["labels", "parent", "id"] + cols)
+    df_all_trees = pd.DataFrame(columns=reserved_names + cols)
     #  Set column type here (useful for continuous vs discrete colorscale)
     for col in cols:
         df_all_trees[col] = df_all_trees[col].astype(df[col].dtype)
