feat(formatters): Add consistent Series formatting for pandas and polars (marimo-team#3232)

eugene · web-flow · commit 6d7c78831110 · 2024-12-19T09:04:38.000-05:00
diff --git a/marimo/_output/formatters/df_formatters.py b/marimo/_output/formatters/df_formatters.py
@@ -43,6 +43,21 @@ def _show_marimo_dataframe(
                     LOGGER.warning("Failed to format DataFrame: %s", e)
                     return ("text/html", df._repr_html_())
 
+            @formatting.opinionated_formatter(pl.Series)
+            def _show_marimo_series(
+                series: pl.Series,
+            ) -> tuple[KnownMimeType, str]:
+                try:
+                    # Table need a column name for operations
+                    if series.name is None or series.name == "":
+                        df = pl.DataFrame({"value": series})
+                    else:
+                        df = series.to_frame()
+                    return table(df, selection=None, pagination=True)._mime_()
+                except Exception as e:
+                    LOGGER.warning("Failed to format Series: %s", e)
+                    return ("text/html", series._repr_html_())
+
 
 class PyArrowFormatter(FormatterFactory):
     @staticmethod
diff --git a/marimo/_output/formatters/pandas_formatters.py b/marimo/_output/formatters/pandas_formatters.py
@@ -3,12 +3,15 @@
 
 from typing import Any
 
+from marimo._loggers import marimo_logger
 from marimo._messaging.mimetypes import KnownMimeType
 from marimo._output.formatters.df_formatters import include_opinionated
 from marimo._output.formatters.formatter_factory import FormatterFactory
 from marimo._output.utils import flatten_string
 from marimo._plugins.ui._impl.table import table
 
+LOGGER = marimo_logger()
+
 
 class PandasFormatter(FormatterFactory):
     @staticmethod
@@ -32,6 +35,21 @@ def _show_marimo_dataframe(
             ) -> tuple[KnownMimeType, str]:
                 return table(df, selection=None, pagination=True)._mime_()
 
+            @formatting.opinionated_formatter(pd.Series)
+            def _show_marimo_series(
+                series: pd.Series[Any],
+            ) -> tuple[KnownMimeType, str]:
+                try:
+                    # Table need a column name for operations
+                    if series.name is None:
+                        series = series.rename("value")
+                    return table(
+                        series.to_frame(), selection=None, pagination=True
+                    )._mime_()
+                except Exception as e:
+                    LOGGER.warning("Failed to format Series: %s", e)
+                    return ("text/html", series._repr_html_())
+
         @formatting.formatter(pd.DataFrame)
         def _show_dataframe(df: pd.DataFrame) -> tuple[KnownMimeType, str]:
             max_rows = pd.get_option("display.max_rows")
diff --git a/marimo/_smoke_tests/series_formatting.py b/marimo/_smoke_tests/series_formatting.py
@@ -0,0 +1,72 @@
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#     "marimo",
+#     "pandas==2.2.3",
+#     "polars==1.17.1",
+# ]
+# ///
+
+import marimo
+
+__generated_with = "0.10.5"
+app = marimo.App(width="medium")
+
+
+@app.cell
+def _():
+    import marimo as mo
+    import pandas as pd
+    import polars as pl
+    return mo, pd, pl
+
+
+@app.cell
+def _(mo):
+    num_range = mo.ui.slider(start=10, step=1, stop=100)
+    num_range
+    return (num_range,)
+
+
+@app.cell
+def _(num_range, pd):
+    # Pandas Series with name
+    pd_named = pd.Series(list(range(num_range.value)), name="numbers")
+    pd_named
+    return (pd_named,)
+
+
+@app.cell
+def _(num_range, pd):
+    # Pandas Series without name
+    pd_unnamed = pd.Series(list(range(num_range.value)))
+    pd_unnamed
+    return (pd_unnamed,)
+
+
+@app.cell
+def _(num_range, pl):
+    # Polars Series with name
+    pl_named = pl.Series("numbers", list(range(num_range.value)))
+    pl_named
+    return (pl_named,)
+
+
+@app.cell
+def _(num_range, pl):
+    # Polars Series without name
+    pl_unnamed = pl.Series(list(range(num_range.value)))
+    pl_unnamed
+    return (pl_unnamed,)
+
+
+@app.cell
+def _(num_range, pl):
+    # Polars Series with empty string name
+    pl_empty_name = pl.Series("", list(range(num_range.value)))
+    pl_empty_name
+    return (pl_empty_name,)
+
+
+if __name__ == "__main__":
+    app.run()
