fix: data table pagination with arrows

frontend/src/components/data-table/__test__/pagination.test.tsx

@@ -0,0 +1,151 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { expect, test } from "vitest";
+import { render } from "@testing-library/react";
+import { PageSelector } from "../pagination";
+import { Functions } from "@/utils/functions";
+
+function getOptions(currentPage: number) {
+  const { container } = render(
+    <PageSelector
+      currentPage={currentPage}
+      totalPages={200}
+      onPageChange={Functions.NOOP}
+    />,
+  );
+
+  const options = container.querySelectorAll("option");
+  const optionValues = [...options].map((option) => option.textContent);
+  return optionValues;
+}
+
+test("pagination start / middle / end", () => {
+  expect(getOptions(1)).toMatchInlineSnapshot(`
+    [
+      "1",
+      "2",
+      "3",
+      "4",
+      "5",
+      "6",
+      "7",
+      "8",
+      "9",
+      "10",
+      "...",
+      "96",
+      "97",
+      "98",
+      "99",
+      "100",
+      "101",
+      "102",
+      "103",
+      "104",
+      "105",
+      "...",
+      "191",
+      "192",
+      "193",
+      "194",
+      "195",
+      "196",
+      "197",
+      "198",
+      "199",
+      "200",
+    ]
+  `);
+
+  // all fall in the top/middle/bottom 10
+  expect(getOptions(1)).toEqual(getOptions(10));
+  expect(getOptions(96)).toEqual(getOptions(105));
+  expect(getOptions(191)).toEqual(getOptions(200));
+
+  // Check off by one
+  expect(getOptions(1)).not.toEqual(getOptions(11));
+  expect(getOptions(1)).not.toEqual(getOptions(95));
+  expect(getOptions(1)).not.toEqual(getOptions(106));
+  expect(getOptions(1)).not.toEqual(getOptions(190));
+});
+
+test("pagination lower middle", () => {
+  expect(getOptions(50)).toMatchInlineSnapshot(`
+    [
+      "1",
+      "2",
+      "3",
+      "4",
+      "5",
+      "6",
+      "7",
+      "8",
+      "9",
+      "10",
+      "...",
+      "50",
+      "...",
+      "96",
+      "97",
+      "98",
+      "99",
+      "100",
+      "101",
+      "102",
+      "103",
+      "104",
+      "105",
+      "...",
+      "191",
+      "192",
+      "193",
+      "194",
+      "195",
+      "196",
+      "197",
+      "198",
+      "199",
+      "200",
+    ]
+  `);
+});
+
+test("pagination upper middle", () => {
+  expect(getOptions(150)).toMatchInlineSnapshot(`
+    [
+      "1",
+      "2",
+      "3",
+      "4",
+      "5",
+      "6",
+      "7",
+      "8",
+      "9",
+      "10",
+      "...",
+      "96",
+      "97",
+      "98",
+      "99",
+      "100",
+      "101",
+      "102",
+      "103",
+      "104",
+      "105",
+      "...",
+      "150",
+      "...",
+      "191",
+      "192",
+      "193",
+      "194",
+      "195",
+      "196",
+      "197",
+      "198",
+      "199",
+      "200",
+    ]
+  `);
+});

frontend/src/components/data-table/pagination.tsx

@@ -9,6 +9,7 @@ import {
 
 import { Button } from "@/components/ui/button";
 import { PluralWord } from "@/utils/pluralize";
+import { range } from "lodash-es";
 
 interface DataTablePaginationProps<TData> {
   table: Table<TData>;
@@ -119,34 +120,11 @@ export const DataTablePagination = <TData,>({
         </Button>
         <div className="flex items-center justify-center text-xs font-medium gap-1">
           <span>Page</span>
-          <select
-            className="cursor-pointer border rounded"
-            value={currentPage}
-            data-testid="page-select"
-            onChange={(e) => table.setPageIndex(Number(e.target.value) - 1)}
-          >
-            {/* If this is too large, this can cause the browser to hang. */}
-            {totalPages <= 100 ? (
-              Array.from({ length: totalPages }, (_, i) => renderOption(i))
-            ) : (
-              // Show the first 10 pages, the middle 10 pages, and the last 10 pages.
-              <>
-                {Array.from({ length: 10 }, (_, i) => renderOption(i))}
-                <option disabled={true} value="__1__">
-                  ...
-                </option>
-                {Array.from({ length: 10 }, (_, i) =>
-                  renderOption(Math.floor(totalPages / 2) - 5 + i),
-                )}
-                <option disabled={true} value="__2__">
-                  ...
-                </option>
-                {Array.from({ length: 10 }, (_, i) =>
-                  renderOption(totalPages - 10 + i),
-                )}
-              </>
-            )}
-          </select>
+          <PageSelector
+            currentPage={currentPage}
+            totalPages={totalPages}
+            onPageChange={(page) => table.setPageIndex(page)}
+          />
           <span className="flex-shrink-0">of {prettyNumber(totalPages)}</span>
         </div>
         <Button
@@ -179,3 +157,80 @@ export const DataTablePagination = <TData,>({
 function prettyNumber(value: number): string {
   return new Intl.NumberFormat().format(value);
 }
+
+export const PageSelector = ({
+  currentPage,
+  totalPages,
+  onPageChange,
+}: {
+  currentPage: number;
+  totalPages: number;
+  onPageChange: (page: number) => void;
+}) => {
+  const renderOption = (i: number) => (
+    <option key={i} value={i + 1}>
+      {i + 1}
+    </option>
+  );
+
+  const renderEllipsis = (key: number) => (
+    <option key={`__${key}__`} disabled={true} value={`__${key}__`}>
+      ...
+    </option>
+  );
+
+  const renderPageOptions = () => {
+    /* If this is too large, this can cause the browser to hang. */
+    if (totalPages <= 100) {
+      return range(totalPages).map((i) => renderOption(i));
+    }
+
+    const middle = Math.floor(totalPages / 2);
+
+    // Show the first 10 pages, the middle 10 pages, and the last 10 pages.
+    const firstPages = range(10).map((i) => renderOption(i));
+    const middlePages = range(10).map((i) => renderOption(middle - 5 + i));
+    const lastPages = range(10).map((i) => renderOption(totalPages - 10 + i));
+
+    const result = [
+      ...firstPages,
+      renderEllipsis(1),
+      ...middlePages,
+      renderEllipsis(2),
+      ...lastPages,
+    ];
+
+    if (currentPage > 10 && currentPage <= middle - 5) {
+      result.splice(
+        10,
+        1, // delete the first ellipsis
+        renderEllipsis(1),
+        renderOption(currentPage - 1),
+        renderEllipsis(11),
+      );
+    }
+
+    if (currentPage > middle + 5 && currentPage <= totalPages - 10) {
+      result.splice(
+        -11,
+        1, // delete the first ellipsis
+        renderEllipsis(2),
+        renderOption(currentPage - 1),
+        renderEllipsis(22),
+      );
+    }
+
+    return result;
+  };
+
+  return (
+    <select
+      className="cursor-pointer border rounded"
+      value={currentPage}
+      data-testid="page-select"
+      onChange={(e) => onPageChange(Number(e.target.value) - 1)}
+    >
+      {renderPageOptions()}
+    </select>
+  );
+};