Add column sizing mode to data editor (#4037)

frontend/src/plugins/impl/DataEditorPlugin.tsx

@@ -57,6 +57,7 @@ export const DataEditorPlugin = createPlugin<Edits>("marimo-data-editor", {
           ]),
         )
         .nullish(),
+      columnSizingMode: z.enum(["auto", "fit"]).default("auto"),
     }),
   )
   .withFunctions({})
@@ -70,6 +71,7 @@ export const DataEditorPlugin = createPlugin<Edits>("marimo-data-editor", {
           fieldTypes={props.data.fieldTypes}
           edits={props.value.edits}
           onEdits={props.setValue}
+          columnSizingMode={props.data.columnSizingMode}
         />
       </TooltipProvider>
     );
@@ -138,6 +140,7 @@ const LoadingDataEditor = (props: Props) => {
         );
         props.onEdits((v) => ({ ...v, edits: [...v.edits, ...newEdits] }));
       }}
+      columnSizingMode={props.columnSizingMode}
     />
   );
 };

frontend/src/plugins/impl/data-editor/data-editor.tsx

@@ -1,5 +1,5 @@
 /* Copyright 2024 Marimo. All rights reserved. */
-import React, { useCallback, useMemo, useRef } from "react";
+import React, { useCallback, useEffect, useMemo, useRef } from "react";
 import { AgGridReact } from "ag-grid-react";
 import type {
   ColDef,
@@ -32,6 +32,7 @@ export interface DataEditorProps<T> {
     }>,
   ) => void;
   onAddRows: (newRows: object[]) => void;
+  columnSizingMode: "fit" | "auto";
 }
 
 function cellEditorForDataType(dataType: DataType) {
@@ -67,6 +68,7 @@ const DataEditor: React.FC<DataEditorProps<object>> = ({
   fieldTypes,
   onAddEdits,
   onAddRows,
+  columnSizingMode,
 }) => {
   const { theme } = useTheme();
   const gridRef = useRef<AgGridReact>(null);
@@ -126,9 +128,31 @@ const DataEditor: React.FC<DataEditorProps<object>> = ({
     };
   }, []);
 
-  const onGridReady = useCallback((params: GridReadyEvent) => {
-    params.api.sizeColumnsToFit();
-  }, []);
+  const sizeColumns = useCallback(() => {
+    // Size the columns to either fit the grid or auto-size based on content
+    const api = gridRef.current?.api;
+    if (!api) {
+      return;
+    }
+
+    if (columnSizingMode === "fit") {
+      api.sizeColumnsToFit();
+    } else if (columnSizingMode === "auto") {
+      api.autoSizeAllColumns();
+    }
+  }, [columnSizingMode]);
+
+  const onGridReady = useCallback(
+    (params: GridReadyEvent) => {
+      sizeColumns();
+    },
+    [sizeColumns],
+  );
+
+  useEffect(() => {
+    // Update grid layout when column sizing prop changes
+    sizeColumns();
+  }, [columnSizingMode, sizeColumns]);
 
   const onCellEditingStopped = useCallback(
     (event: CellEditingStoppedEvent) => {

marimo/_plugins/ui/_impl/data_editor.py

@@ -9,6 +9,7 @@
     Any,
     Callable,
     Final,
+    Literal,
     Optional,
     TypedDict,
     Union,
@@ -121,6 +122,8 @@ class data_editor(
             Can be a Pandas dataframe, a list of dicts, or a dict of lists.
         label (str): Markdown label for the element.
         on_change (Optional[Callable]): Optional callback to run when this element's value changes.
+        column_sizing_mode (Literal["auto", "fit"]): The column sizing mode for the table.
+            `auto` will size columns based on the content, `fit` will size columns to fit the view.
     """
 
     _name: Final[str] = "marimo-data-editor"
@@ -140,6 +143,7 @@ def __init__(
                 None,
             ]
         ] = None,
+        column_sizing_mode: Literal["auto", "fit"] = "auto",
     ) -> None:
         validate_page_size(page_size)
         table_manager = get_table_manager(data)
@@ -163,6 +167,7 @@ def __init__(
                 "field-types": field_types or None,
                 "pagination": pagination,
                 "page-size": page_size,
+                "column-sizing-mode": column_sizing_mode,
             },
             on_change=on_change,
         )

tests/_plugins/ui/_impl/test_data_editor.py

@@ -30,6 +30,7 @@ def test_data_editor_initialization():
     assert editor._edits == {"edits": []}
     assert editor._component_args["pagination"] is True
     assert editor._component_args["page-size"] == 50
+    assert editor._component_args["column-sizing-mode"] == "auto"
 
 
 @pytest.mark.skipif(