@marimo-team/frontend 0.25.2-dev56 → 0.25.2-dev62

package/src/components/data-table/schemas.ts

@@ -3,7 +3,7 @@
 import z from "zod";
 import { rpc } from "@/plugins/core/rpc";
 
-export type DownloadFormat = "csv" | "json" | "parquet" | "tsv";
+export type DownloadFormat = "csv" | "json" | "parquet" | "tsv" | "geojson";
 
 /**
  * Per-request export options. Each field is optional. A missing field keeps
@@ -18,6 +18,7 @@ export interface DownloadAsOptions {
 export interface DownloadAsRequest {
   format: DownloadFormat;
   options?: DownloadAsOptions;
+  geometry_column?: string | null;
 }
 
 export type DownloadAsArgs = (req: DownloadAsRequest) => Promise<{
@@ -25,26 +26,101 @@ export type DownloadAsArgs = (req: DownloadAsRequest) => Promise<{
   filename: string;
   error?: string | null;
   missing_packages?: string[] | null;
+  code?:
+    | "geometry_required"
+    | "invalid_geometry"
+    | "invalid_metadata"
+    | "unsupported_representation"
+    | "unsupported_version"
+    | "missing_packages"
+    | "missing_crs"
+    | "conversion_failed"
+    | null;
+  column?: string | null;
 }>;
 
+const DownloadRequestSchema = z.object({
+  geometry_column: z.string().nullish(),
+  options: z
+    .object({
+      separator: z.string().optional(),
+      encoding: z.string().optional(),
+      ensure_ascii: z.boolean().optional(),
+    })
+    .optional(),
+});
+
+const DownloadResponseSchema = z.object({
+  url: z.string(),
+  filename: z.string(),
+  error: z.string().nullish(),
+  missing_packages: z.array(z.string()).nullish(),
+  code: z
+    .enum([
+      "geometry_required",
+      "invalid_geometry",
+      "invalid_metadata",
+      "unsupported_representation",
+      "unsupported_version",
+      "missing_packages",
+      "conversion_failed",
+    ])
+    .nullish(),
+  column: z.string().nullish(),
+});
+
 export const DownloadAsSchema = rpc
   .input(
-    z.object({
+    DownloadRequestSchema.extend({
       format: z.enum(["csv", "json", "parquet", "tsv"]),
-      options: z
-        .object({
-          separator: z.string().optional(),
-          encoding: z.string().optional(),
-          ensure_ascii: z.boolean().optional(),
-        })
-        .optional(),
     }),
   )
+  .output(DownloadResponseSchema);
+
+export const DownloadGeoJSONSchema = rpc
+  .input(DownloadRequestSchema.extend({ format: z.literal("geojson") }))
   .output(
-    z.object({
-      url: z.string(),
-      filename: z.string(),
-      error: z.string().nullish(),
-      missing_packages: z.array(z.string()).nullish(),
+    DownloadResponseSchema.extend({
+      code: z.union([
+        DownloadResponseSchema.shape.code,
+        z.literal("missing_crs"),
+      ]),
     }),
   );
+
+export type DownloadAsFunction = (
+  req: z.infer<typeof DownloadAsSchema.input>,
+) => Promise<z.infer<typeof DownloadAsSchema.output>>;
+
+export type DownloadGeoJSON = (
+  req: z.infer<typeof DownloadGeoJSONSchema.input>,
+) => Promise<z.infer<typeof DownloadGeoJSONSchema.output>>;
+
+const ExportMetadataSchema = z.object({
+  geometry_columns: z.array(
+    z.object({
+      name: z.string(),
+      encoding: z.enum(["objects", "wkb", "wkt", "other"]),
+      crs: z.union([z.string(), z.record(z.string(), z.unknown())]).nullable(),
+    }),
+  ),
+  primary_geometry_column: z.string().nullable(),
+  default_geometry_column: z.string().nullable(),
+  formats: z.record(
+    z.string(),
+    z.object({
+      available: z.boolean(),
+      reason: z.string().nullable(),
+      missing_packages: z.array(z.string()),
+    }),
+  ),
+});
+
+export type ExportMetadata = z.infer<typeof ExportMetadataSchema>;
+export type GetExportMetadata = (
+  opts: Record<string, never>,
+) => Promise<ExportMetadata>;
+
+export const GetExportMetadataSchema = rpc
+  .input(z.object({}))
+  .output(ExportMetadataSchema);

package/src/plugins/impl/DataTablePlugin.tsx

@@ -36,7 +36,12 @@ import { usePanelOwnership } from "@/components/data-table/hooks/use-panel-owner
 import { LoadingTable } from "@/components/data-table/loading-table";
 import {
   type DownloadAsArgs,
+  type DownloadAsFunction,
   DownloadAsSchema,
+  type DownloadGeoJSON,
+  DownloadGeoJSONSchema,
+  type GetExportMetadata,
+  GetExportMetadataSchema,
 } from "@/components/data-table/schemas";
 import { TableExplorerPanel } from "@/components/data-table/table-explorer-panel/table-explorer-panel";
 import {
@@ -209,7 +214,9 @@ interface Data<T> {
 
 // oxlint-disable-next-line typescript/consistent-type-definitions
 type DataTableFunctions = {
-  download_as: DownloadAsArgs;
+  download_as: DownloadAsFunction;
+  download_geojson: DownloadGeoJSON;
+  get_export_metadata: GetExportMetadata;
   get_column_summaries: <T>(opts: {}) => Promise<ColumnSummaries<T>>;
   search: <T>(req: {
     sort?: {
@@ -302,6 +309,8 @@ export const DataTablePlugin = createPlugin<S>("marimo-table")
   )
   .withFunctions<DataTableFunctions>({
     download_as: DownloadAsSchema,
+    download_geojson: DownloadGeoJSONSchema,
+    get_export_metadata: GetExportMetadataSchema,
     get_column_summaries: rpc.input(z.looseObject({})).output(
       z.object({
         data: z.union([z.string(), z.array(z.looseObject({}))]).nullable(),
@@ -839,7 +848,9 @@ const DataTableComponent = ({
   fieldTypes,
   paginationState,
   setPaginationState,
-  download_as: downloadAs,
+  download_as,
+  download_geojson,
+  get_export_metadata: getExportMetadata,
   columnSummaries,
   className,
   setValue,
@@ -880,6 +891,15 @@ const DataTableComponent = ({
     sizeBytesIsLoading?: boolean;
   }): JSX.Element => {
   const id = useId();
+  const downloadAs = useCallback<DownloadAsArgs>(
+    (request) => {
+      const { format, ...options } = request;
+      return format === "geojson"
+        ? download_geojson({ format, ...options })
+        : download_as({ format, ...options });
+    },
+    [download_as, download_geojson],
+  );
   const { locale } = useLocale();
   const [viewedRowIdx, setViewedRowIdx] = useState(0);
   const { isPanelOpen, isAnyPanelOpen, togglePanel, panelType, setPanelType } =
@@ -1168,6 +1188,7 @@ const DataTableComponent = ({
             hoverTemplate={hoverTemplate}
             cellHoverTexts={cellHoverTexts}
             downloadAs={showDownload ? downloadAs : undefined}
+            getExportMetadata={showDownload ? getExportMetadata : undefined}
             showSearch={showSearch}
             searchQuery={searchQuery}
             onSearchQueryChange={setSearchQuery}

package/src/plugins/impl/__tests__/DataTablePlugin.test.tsx

@@ -4,11 +4,17 @@ import { Tooltip } from "radix-ui";
 
 const TooltipProvider = Tooltip.Provider;
 
-import { act, render, screen, waitFor } from "@testing-library/react";
+import {
+  act,
+  fireEvent,
+  render,
+  screen,
+  waitFor,
+} from "@testing-library/react";
 import { Provider } from "jotai";
 import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
 import { SetupMocks } from "@/__mocks__/common";
-import type { DownloadAsArgs } from "@/components/data-table/schemas";
+import type { DownloadAsFunction } from "@/components/data-table/schemas";
 import type { FieldTypesWithExternalType } from "@/components/data-table/types";
 import { store } from "@/core/state/jotai";
 import {
@@ -19,6 +25,10 @@ import {
 
 // Default to normal (non-static) mode; individual tests flip this on.
 const mockIsStatic = vi.fn().mockReturnValue(false);
+vi.mock("@/utils/download", () => ({
+  downloadByURL: vi.fn(),
+  withLoadingToast: vi.fn(async (_title, callback) => callback()),
+}));
 vi.mock("@/core/static/static-state", async (importOriginal) => {
   const actual =
     await importOriginal<typeof import("@/core/static/static-state")>();
@@ -30,9 +40,15 @@ beforeAll(() => {
 });
 
 describe("LoadingDataTableComponent", () => {
-  it("keeps data columns when inferred types include a row header", async () => {
+  it("keeps data columns and routes each export to its RPC", async () => {
     const host = document.createElement("div");
     const data = JSON.stringify([{ index: "first", value: 1 }]);
+    const downloadAs = vi
+      .fn()
+      .mockResolvedValue({ url: "export", filename: "table.csv" });
+    const downloadGeoJSON = vi
+      .fn()
+      .mockResolvedValue({ url: "export", filename: "table.geojson" });
 
     render(
       <Provider store={store}>
@@ -43,7 +59,7 @@ describe("LoadingDataTableComponent", () => {
             pagination={false}
             pageSize={10}
             selection={null}
-            showDownload={false}
+            showDownload={true}
             showFilters={false}
             showColumnSummaries={false}
             showDataTypes={false}
@@ -68,7 +84,22 @@ describe("LoadingDataTableComponent", () => {
               cell_styles: null,
               cell_hover_texts: null,
             })}
-            download_as={vi.fn() as DownloadAsArgs}
+            download_as={downloadAs}
+            download_geojson={downloadGeoJSON}
+            get_export_metadata={vi.fn().mockResolvedValue({
+              geometry_columns: [
+                { name: "value", encoding: "objects", crs: "EPSG:4326" },
+              ],
+              primary_geometry_column: "value",
+              default_geometry_column: "value",
+              formats: {
+                geojson: {
+                  available: true,
+                  reason: null,
+                  missing_packages: [],
+                },
+              },
+            })}
             get_column_summaries={vi.fn()}
             get_data_url={vi.fn() as GetDataUrl}
             get_row_ids={vi.fn() as GetRowIds}
@@ -83,6 +114,26 @@ describe("LoadingDataTableComponent", () => {
         screen.getByRole("columnheader", { name: "value" }),
       ).toBeInTheDocument();
     });
+    fireEvent.click(screen.getByTestId("export-button"));
+    await screen.findByRole("dialog", { name: "Export table" });
+    await waitFor(() =>
+      expect(
+        screen.getByRole("button", { name: "Download GeoJSON" }),
+      ).toBeEnabled(),
+    );
+    fireEvent.click(screen.getByRole("button", { name: "Download GeoJSON" }));
+    await waitFor(() =>
+      expect(downloadGeoJSON).toHaveBeenCalledWith({
+        format: "geojson",
+        geometry_column: "value",
+      }),
+    );
+    expect(downloadAs).not.toHaveBeenCalled();
+    fireEvent.click(screen.getByRole("button", { name: "Download CSV" }));
+    await waitFor(() =>
+      expect(downloadAs).toHaveBeenCalledWith({ format: "csv" }),
+    );
+    expect(downloadGeoJSON).toHaveBeenCalledTimes(1);
   });
 
   it("does not log duplicate keys when a pivot becomes empty", async () => {
@@ -122,7 +173,9 @@ describe("LoadingDataTableComponent", () => {
       showSearch: false,
       value: [] as (number | string | { rowId: string; columnName?: string })[],
       setValue: vi.fn(),
-      download_as: vi.fn() as DownloadAsArgs,
+      download_as: vi.fn() as DownloadAsFunction,
+      download_geojson: vi.fn(),
+      get_export_metadata: vi.fn(),
       get_column_summaries: vi.fn(),
       get_data_url: vi.fn() as GetDataUrl,
       get_row_ids: vi.fn() as GetRowIds,
@@ -263,7 +316,9 @@ describe("LoadingDataTableComponent", () => {
       showSearch: true,
       value: [] as (number | string | { rowId: string; columnName?: string })[],
       setValue,
-      download_as: vi.fn() as DownloadAsArgs,
+      download_as: vi.fn() as DownloadAsFunction,
+      download_geojson: vi.fn(),
+      get_export_metadata: vi.fn(),
       get_column_summaries: vi.fn().mockResolvedValue({
         data: null,
         stats: {},
@@ -379,7 +434,9 @@ describe("static notebook control suppression", () => {
         cell_styles: null,
         cell_hover_texts: null,
       }),
-      download_as: vi.fn() as DownloadAsArgs,
+      download_as: vi.fn() as DownloadAsFunction,
+      download_geojson: vi.fn(),
+      get_export_metadata: vi.fn(),
       get_column_summaries: vi.fn().mockResolvedValue({
         data: null,
         stats: {},

package/src/plugins/impl/data-frames/DataFramePlugin.tsx

@@ -12,8 +12,12 @@ import {
 } from "react";
 import { z } from "zod";
 import {
-  type DownloadAsArgs,
+  type DownloadAsFunction,
   DownloadAsSchema,
+  type DownloadGeoJSON,
+  DownloadGeoJSONSchema,
+  type GetExportMetadata,
+  GetExportMetadataSchema,
 } from "@/components/data-table/schemas";
 import type { FieldTypesWithExternalType } from "@/components/data-table/types";
 import { ReadonlyCode } from "@/components/editor/code/readonly-python-code";
@@ -82,7 +86,9 @@ type PluginFunctions = {
     data: TableData<T>;
     total_rows: number;
   }>;
-  download_as: DownloadAsArgs;
+  download_as: DownloadAsFunction;
+  download_geojson: DownloadGeoJSON;
+  get_export_metadata: GetExportMetadata;
   get_size_bytes: (opts: Record<string, never>) => Promise<{
     size_bytes?: number | null;
   }>;
@@ -153,6 +159,8 @@ export const DataFramePlugin = createPlugin<S>("marimo-dataframe")
         }),
       ),
     download_as: DownloadAsSchema,
+    download_geojson: DownloadGeoJSONSchema,
+    get_export_metadata: GetExportMetadataSchema,
     get_size_bytes: rpc
       .input(z.object({}))
       .output(z.object({ size_bytes: z.number().nullish() })),
@@ -174,7 +182,7 @@ interface DataTableProps extends Data, PluginFunctions {
   setValue: (value: S) => void;
   host: HTMLElement;
   showDownload: boolean;
-  download_as: DownloadAsArgs;
+  download_as: DownloadAsFunction;
 }
 
 const EMPTY: Transformations = {
@@ -194,6 +202,8 @@ export const DataFrameComponent = memo(
     get_column_values,
     search,
     download_as,
+    download_geojson,
+    get_export_metadata,
     get_size_bytes,
     host,
   }: DataTableProps): JSX.Element => {
@@ -337,6 +347,8 @@ export const DataFrameComponent = memo(
           rowHeaders={row_headers || Arrays.EMPTY}
           showDownload={showDownload}
           download_as={download_as}
+          download_geojson={download_geojson}
+          get_export_metadata={get_export_metadata}
           get_size_bytes={get_size_bytes}
           showSearch={false}
           showFilters={false}

package/src/stories/dataframe.stories.tsx

@@ -42,6 +42,13 @@ export const DataFrame: StoryObj = {
         host={document.body}
         showDownload={false}
         download_as={async () => ({ url: "", filename: "" })}
+        download_geojson={async () => ({ url: "", filename: "" })}
+        get_export_metadata={async () => ({
+          geometry_columns: [],
+          primary_geometry_column: null,
+          default_geometry_column: null,
+          formats: {},
+        })}
         get_size_bytes={async () => ({ size_bytes: null })}
         lazy={false}
       />