Merge pull request #4914 from MarcoGorelli/fix-timeline

fix: px.timeline was raising when x_start and/or x_end were already datetime for Polars / PyArrow

Author: emilykl

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

@@ -2118,16 +2118,21 @@ def process_dataframe_timeline(args):
     if args["x_start"] is None or args["x_end"] is None:
         raise ValueError("Both x_start and x_end are required")
 
-    try:
-        df: nw.DataFrame = args["data_frame"]
-        df = df.with_columns(
-            nw.col(args["x_start"]).str.to_datetime().alias(args["x_start"]),
-            nw.col(args["x_end"]).str.to_datetime().alias(args["x_end"]),
-        )
-    except Exception:
-        raise TypeError(
-            "Both x_start and x_end must refer to data convertible to datetimes."
-        )
+    df: nw.DataFrame = args["data_frame"]
+    schema = df.schema
+    to_convert_to_datetime = [
+        col
+        for col in [args["x_start"], args["x_end"]]
+        if schema[col] != nw.Datetime and schema[col] != nw.Date
+    ]
+
+    if to_convert_to_datetime:
+        try:
+            df = df.with_columns(nw.col(to_convert_to_datetime).str.to_datetime())
+        except Exception as exc:
+            raise TypeError(
+                "Both x_start and x_end must refer to data convertible to datetimes."
+            ) from exc
 
     # note that we are not adding any columns to the data frame here, so no risk of overwrite
     args["data_frame"] = df.with_columns(

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

@@ -561,3 +561,29 @@ def test_timeline(constructor):
     msg = "Both x_start and x_end must refer to data convertible to datetimes."
     with pytest.raises(TypeError, match=msg):
         px.timeline(df, x_start="Start", x_end=["a", "b", "c"], y="Task", color="Task")
+
+
+@pytest.mark.parametrize(
+    "datetime_columns",
+    [
+        ["Start"],
+        ["Start", "Finish"],
+        ["Finish"],
+    ],
+)
+def test_timeline_cols_already_temporal(constructor, datetime_columns):
+    # https://github.com/plotly/plotly.py/issues/4913
+    data = {
+        "Task": ["Job A", "Job B", "Job C"],
+        "Start": ["2009-01-01", "2009-03-05", "2009-02-20"],
+        "Finish": ["2009-02-28", "2009-04-15", "2009-05-30"],
+    }
+    df = (
+        nw.from_native(constructor(data))
+        .with_columns(nw.col(datetime_columns).str.to_datetime(format="%Y-%m-%d"))
+        .to_native()
+    )
+    fig = px.timeline(df, x_start="Start", x_end="Finish", y="Task", color="Task")
+    assert len(fig.data) == 3
+    assert fig.layout.xaxis.type == "date"
+    assert fig.layout.xaxis.title.text is None