fix

Author: mscolnick

marimo/_plugins/ui/_impl/altair_chart.py

@@ -15,7 +15,6 @@
 )
 
 import narwhals.stable.v2 as nw
-from narwhals.dependencies import is_narwhals_lazyframe
 from narwhals.typing import IntoDataFrame, IntoLazyFrame
 
 from marimo import _loggers
@@ -55,7 +54,9 @@
 RowOrientedData = list[dict[str, Any]]
 ColumnOrientedData = dict[str, list[Any]]
 
-ChartDataType = Union[IntoDataFrame, RowOrientedData, ColumnOrientedData]
+ChartDataType = Union[
+    IntoDataFrame, IntoLazyFrame, RowOrientedData, ColumnOrientedData
+]
 
 # Union of all possible chart types
 AltairChartType: TypeAlias = "altair.vegalite.v5.api.ChartType"
@@ -98,12 +99,8 @@ def _using_vegafusion() -> bool:
 
 def _filter_dataframe(
     native_df: Union[IntoDataFrame, IntoLazyFrame], selection: ChartSelection
-) -> IntoDataFrame:
-    df = nw.from_native(native_df).lazy()
-    if is_narwhals_lazyframe(df):
-        raise ValueError(
-            "Lazyframes are not supported for filtering. Run `df.collect()` before filtering."
-        )
+) -> Union[IntoDataFrame, IntoLazyFrame]:
+    df = nw.from_native(native_df)
     if not isinstance(selection, dict):
         raise TypeError("Input 'selection' must be a dictionary")
 
@@ -187,14 +184,9 @@ def _coerce_value(value: Any, dtype: Any) -> Any:
         if nw.Datetime == dtype and isinstance(dtype, nw.Datetime):
             if isinstance(value, str):
                 res = datetime.datetime.fromisoformat(value)
-                # If dtype has no timezone, shift by local timezone offset
-                if dtype.time_zone is None:
-                    local_tz = datetime.datetime.now().astimezone().tzinfo
-                    LOGGER.warning(
-                        f"Datetime was given with a timezone when not expected. "
-                        f"Shifting by local timezone offset {local_tz}."
-                    )
-                    return res.astimezone(local_tz).replace(tzinfo=None)
+                # If dtype has no timezone, but value has timezone, remove timezone without shifting
+                if dtype.time_zone is None and res.tzinfo is not None:
+                    return res.replace(tzinfo=None)
                 return res
 
             # Value is milliseconds since epoch

tests/_plugins/ui/_impl/test_altair_chart.py

@@ -665,19 +665,19 @@ def test_parse_spec_pandas() -> None:
     chart = alt.Chart(data).mark_point().encode(x="values:Q")
     spec = _parse_spec(chart)
     # Replace data.url with a placeholder
-    spec["data"]["url"] = "_placeholder_"
+    spec["data"] = {"url": "_placeholder_", "format": spec["data"]["format"]}
     snapshot("parse_spec_pandas.txt", json.dumps(spec, indent=2))
 
 
 @pytest.mark.skipif(not HAS_DEPS, reason="optional dependencies not installed")
 def test_parse_spec_narwhal() -> None:
     import altair as alt
 
-    data = nw.from_native(pd.DataFrame({"values": [1, 2, 3]}))
+    data = pd.DataFrame({"values": [1, 2, 3]})
     chart = alt.Chart(data).mark_point().encode(x="values:Q")
     spec = _parse_spec(chart)
     # Replace data.url with a placeholder
-    spec["data"]["url"] = "_placeholder_"
+    spec["data"] = {"url": "_placeholder_", "format": spec["data"]["format"]}
     snapshot("parse_spec_narwhal.txt", json.dumps(spec, indent=2))
 
 
@@ -690,7 +690,7 @@ def test_parse_spec_polars() -> None:
     chart = alt.Chart(data).mark_point().encode(x="values:Q")
     spec = _parse_spec(chart)
     # Replace data.url with a placeholder
-    spec["data"]["url"] = "_placeholder_"
+    spec["data"] = {"url": "_placeholder_", "format": spec["data"]["format"]}
     snapshot("parse_spec_polars.txt", json.dumps(spec, indent=2))