Update vizro-core/examples/visual-vocabulary/custom_charts.py

Co-authored-by: Petar Pejovic <[email protected]>

vizro-core/examples/visual-vocabulary/custom_charts.py

@@ -336,25 +336,16 @@ def lollipop(data_frame: pd.DataFrame, **kwargs):
     fig = px.scatter(data_frame, **kwargs)
 
     # Enable for both orientations
-    orientation = fig.data[0].orientation
-    x_array = fig.data[0]["x"]
-    y_array = fig.data[0]["y"]
-
-    for i in range(len(data_frame)):
-        fig.add_trace(
-            go.Scatter(
-                x=[0, x_array[i]] if orientation == "h" else [x_array[i], x_array[i]],
-                y=[y_array[i], y_array[i]] if orientation == "h" else [0, y_array[i]],
-                mode="lines",
-            )
-        )
+    is_horizontal = fig.data[0].orientation == "h"
 
-    if orientation == "h":
-        yaxis_showgrid = False
-        xaxis_showgrid = True
-    else:
-        yaxis_showgrid = True
-        xaxis_showgrid = False
+    x_coords = [[0, x] if is_horizontal else [x, x] for x in fig.data[0]["x"]]
+    y_coords = [[y, y] if is_horizontal else [0, y] for y in fig.data[0]["y"]]
+
+    for x, y in zip(x_coords, y_coords):
+        fig.add_trace(go.Scatter(x=x, y=y, mode="lines"))
+
+    xaxis_showgrid = is_horizontal
+    yaxis_showgrid = not is_horizontal
 
     fig.update_traces(
         marker_size=12,