From d469706cdf63d1a2c00462636eeba4539e8e1190 Mon Sep 17 00:00:00 2001 From: Ken McGrady Date: Sat, 18 Oct 2025 12:40:38 -0700 Subject: [PATCH] Break up AppNode into separate files --- frontend/lib/src/AppNode.test.ts | 1212 ----------------- frontend/lib/src/AppNode.ts | 929 +------------ .../lib/src/render-tree/AppNode.interface.ts | 122 ++ frontend/lib/src/render-tree/AppRoot.test.ts | 664 +++++++++ frontend/lib/src/render-tree/AppRoot.ts | 426 ++++++ .../lib/src/render-tree/BlockNode.test.ts | 74 + frontend/lib/src/render-tree/BlockNode.ts | 196 +++ .../lib/src/render-tree/ElementNode.test.ts | 421 ++++++ frontend/lib/src/render-tree/ElementNode.ts | 273 ++++ frontend/lib/src/render-tree/test-utils.ts | 173 +++ frontend/vitest.config.ts | 6 +- 11 files changed, 2358 insertions(+), 2138 deletions(-) delete mode 100644 frontend/lib/src/AppNode.test.ts create mode 100644 frontend/lib/src/render-tree/AppNode.interface.ts create mode 100644 frontend/lib/src/render-tree/AppRoot.test.ts create mode 100644 frontend/lib/src/render-tree/AppRoot.ts create mode 100644 frontend/lib/src/render-tree/BlockNode.test.ts create mode 100644 frontend/lib/src/render-tree/BlockNode.ts create mode 100644 frontend/lib/src/render-tree/ElementNode.test.ts create mode 100644 frontend/lib/src/render-tree/ElementNode.ts create mode 100644 frontend/lib/src/render-tree/test-utils.ts diff --git a/frontend/lib/src/AppNode.test.ts b/frontend/lib/src/AppNode.test.ts deleted file mode 100644 index 003cbaa405b..00000000000 --- a/frontend/lib/src/AppNode.test.ts +++ /dev/null @@ -1,1212 +0,0 @@ -/** - * Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2025) - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -import { Writer } from "protobufjs" -import { MockInstance } from "vitest" - -import { - ArrowNamedDataSet, - Block as BlockProto, - Delta as DeltaProto, - Element, - ForwardMsgMetadata, - IArrowVegaLiteChart, - Logo as LogoProto, -} from "@streamlit/protobuf" - -import { isNullOrUndefined } from "~lib/util/utils" - -import { AppNode, AppRoot, BlockNode, ElementNode } from "./AppNode" -import { UNICODE } from "./mocks/arrow" - -const NO_SCRIPT_RUN_ID = "NO_SCRIPT_RUN_ID" -const FAKE_SCRIPT_HASH = "fake_script_hash" -// prettier-ignore -const BLOCK = block([ - text("1"), - block([ - text("2"), - ]), -]) - -// Initialize new AppRoot with a main block node and three child block nodes - sidebar, events and bottom. -const ROOT = new AppRoot( - FAKE_SCRIPT_HASH, - new BlockNode(FAKE_SCRIPT_HASH, [ - BLOCK, - new BlockNode(FAKE_SCRIPT_HASH), - new BlockNode(FAKE_SCRIPT_HASH), - new BlockNode(FAKE_SCRIPT_HASH), - ]) -) - -describe("AppNode.getIn", () => { - it("handles shallow paths", () => { - const node = BLOCK.getIn([0]) - expect(node).toBeTextNode("1") - }) - - it("handles deep paths", () => { - const node = BLOCK.getIn([1, 0]) - expect(node).toBeTextNode("2") - }) - - it("returns undefined for invalid paths", () => { - const node = BLOCK.getIn([2, 3, 4]) - expect(node).toBeUndefined() - }) -}) - -describe("AppNode.setIn", () => { - it("handles shallow paths", () => { - const newBlock = BLOCK.setIn([0], text("new"), NO_SCRIPT_RUN_ID) - expect(newBlock.getIn([0])).toBeTextNode("new") - - // Check BLOCK..newBlock diff is as expected. - expect(newBlock).not.toStrictEqual(BLOCK) - expect(newBlock.getIn([1])).toStrictEqual(BLOCK.getIn([1])) - }) - - it("handles deep paths", () => { - const newBlock = BLOCK.setIn([1, 1], text("new"), NO_SCRIPT_RUN_ID) - expect(newBlock.getIn([1, 1])).toBeTextNode("new") - - // Check BLOCK..newBlock diff is as expected - expect(newBlock).not.toStrictEqual(BLOCK) - expect(newBlock.getIn([0])).toStrictEqual(BLOCK.getIn([0])) - expect(newBlock.getIn([1])).not.toStrictEqual(BLOCK.getIn([1])) - expect(newBlock.getIn([1, 0])).toStrictEqual(BLOCK.getIn([1, 0])) - expect(newBlock.getIn([1, 1])).not.toStrictEqual(BLOCK.getIn([1, 1])) - }) - - it("throws an error for invalid paths", () => { - expect(() => BLOCK.setIn([1, 2], text("new"), NO_SCRIPT_RUN_ID)).toThrow( - "Bad 'setIn' index 2 (should be between [0, 1])" - ) - }) -}) - -describe("ElementNode.quiverElement", () => { - it("returns a quiverElement (arrowTable)", () => { - const node = arrowTable() - const q = node.quiverElement - expect(q.columnNames).toEqual([["", "c1", "c2"]]) - expect(q.getCell(0, 0).content).toEqual("i1") - }) - - it("returns a quiverElement (arrowDataFrame)", () => { - const node = arrowDataFrame() - const q = node.quiverElement - expect(q.columnNames).toEqual([["", "c1", "c2"]]) - expect(q.getCell(0, 0).content).toEqual("i1") - }) - - it("does not recompute its value (arrowTable)", () => { - // accessing `quiverElement` twice should return the same instance. - const node = arrowTable() - expect(node.quiverElement).toStrictEqual(node.quiverElement) - }) - - it("does not recompute its value (arrowDataFrame)", () => { - // accessing `quiverElement` twice should return the same instance. - const node = arrowDataFrame() - expect(node.quiverElement).toStrictEqual(node.quiverElement) - }) - - it("throws an error for other element types", () => { - const node = text("foo") - expect(() => node.quiverElement).toThrow( - "elementType 'text' is not a valid Quiver element!" - ) - }) -}) - -describe("ElementNode.vegaLiteChartElement", () => { - it("returns a vegaLiteChartElement (data)", () => { - const MOCK_VEGA_LITE_CHART = { - spec: JSON.stringify({ - mark: "circle", - encoding: { - x: { field: "a", type: "quantitative" }, - y: { field: "b", type: "quantitative" }, - size: { field: "c", type: "quantitative" }, - color: { field: "c", type: "quantitative" }, - }, - }), - data: { data: UNICODE }, - datasets: [], - useContainerWidth: true, - } - const node = arrowVegaLiteChart(MOCK_VEGA_LITE_CHART) - const element = node.vegaLiteChartElement - - // spec - expect(element.spec).toEqual(MOCK_VEGA_LITE_CHART.spec) - - // data - expect(element.data?.columnNames).toEqual([["", "c1", "c2"]]) - expect(element.data?.getCell(0, 0).content).toEqual("i1") - - // datasets - expect(element.datasets.length).toEqual(0) - - // use container width - expect(element.useContainerWidth).toEqual( - MOCK_VEGA_LITE_CHART.useContainerWidth - ) - }) - - it("returns a vegaLiteChartElement (datasets)", () => { - const MOCK_VEGA_LITE_CHART = { - spec: JSON.stringify({ - mark: "circle", - encoding: { - x: { field: "a", type: "quantitative" }, - y: { field: "b", type: "quantitative" }, - size: { field: "c", type: "quantitative" }, - color: { field: "c", type: "quantitative" }, - }, - }), - data: null, - datasets: [{ hasName: true, name: "foo", data: { data: UNICODE } }], - useContainerWidth: true, - } - const node = arrowVegaLiteChart(MOCK_VEGA_LITE_CHART) - const element = node.vegaLiteChartElement - - // spec - expect(element.spec).toEqual(MOCK_VEGA_LITE_CHART.spec) - - // data - expect(element.data).toEqual(null) - - // datasets - expect(element.datasets[0].hasName).toEqual( - MOCK_VEGA_LITE_CHART.datasets[0].hasName - ) - expect(element.datasets[0].name).toEqual( - MOCK_VEGA_LITE_CHART.datasets[0].name - ) - expect(element.datasets[0].data.columnNames).toEqual([["", "c1", "c2"]]) - expect(element.datasets[0].data.getCell(0, 0).content).toEqual("i1") - - // use container width - expect(element.useContainerWidth).toEqual( - MOCK_VEGA_LITE_CHART.useContainerWidth - ) - }) - - it("does not recompute its value", () => { - const MOCK_VEGA_LITE_CHART = { - spec: JSON.stringify({ - mark: "circle", - encoding: { - x: { field: "a", type: "quantitative" }, - y: { field: "b", type: "quantitative" }, - size: { field: "c", type: "quantitative" }, - color: { field: "c", type: "quantitative" }, - }, - }), - data: { data: UNICODE }, - datasets: [], - useContainerWidth: true, - } - // accessing `vegaLiteChartElement` twice should return the same instance. - const node = arrowVegaLiteChart(MOCK_VEGA_LITE_CHART) - expect(node.vegaLiteChartElement).toStrictEqual(node.vegaLiteChartElement) - }) - - it("throws an error for other element types", () => { - const node = text("foo") - expect(() => node.vegaLiteChartElement).toThrow( - "elementType 'text' is not a valid VegaLiteChartElement!" - ) - }) -}) - -describe("ElementNode.arrowAddRows", () => { - const MOCK_UNNAMED_DATASET = { - hasName: false, - name: "", - data: { data: UNICODE }, - } as ArrowNamedDataSet - const MOCK_NAMED_DATASET = { - hasName: true, - name: "foo", - data: { data: UNICODE }, - } as ArrowNamedDataSet - const MOCK_ANOTHER_NAMED_DATASET = { - hasName: true, - name: "bar", - data: { data: UNICODE }, - } as ArrowNamedDataSet - - describe("arrowTable", () => { - test("addRows can be called with an unnamed dataset", () => { - const node = arrowTable() - const newNode = node.arrowAddRows(MOCK_UNNAMED_DATASET, NO_SCRIPT_RUN_ID) - const q = newNode.quiverElement - - expect(q.columnNames).toEqual([["", "c1", "c2"]]) - expect(q.dimensions.numDataRows).toEqual(4) - expect(q.getCell(0, 0).content).toEqual("i1") - expect(q.getCell(2, 0).content).toEqual("i1") - expect(q.getCell(0, 1).content).toEqual("foo") - expect(q.getCell(2, 1).content).toEqual("foo") - }) - - test("addRows throws an error when called with a named dataset", () => { - const node = arrowTable() - expect(() => - node.arrowAddRows(MOCK_NAMED_DATASET, NO_SCRIPT_RUN_ID) - ).toThrow( - "Add rows cannot be used with a named dataset for this element." - ) - }) - }) - - describe("arrowDataFrame", () => { - test("addRows can be called with an unnamed dataset", () => { - const node = arrowDataFrame() - const newNode = node.arrowAddRows(MOCK_UNNAMED_DATASET, NO_SCRIPT_RUN_ID) - const q = newNode.quiverElement - - expect(q.columnNames).toEqual([["", "c1", "c2"]]) - expect(q.dimensions.numDataRows).toEqual(4) - expect(q.getCell(0, 0).content).toEqual("i1") - expect(q.getCell(2, 0).content).toEqual("i1") - expect(q.getCell(0, 1).content).toEqual("foo") - expect(q.getCell(2, 1).content).toEqual("foo") - }) - - test("addRows throws an error when called with a named dataset", () => { - const node = arrowDataFrame() - expect(() => - node.arrowAddRows(MOCK_NAMED_DATASET, NO_SCRIPT_RUN_ID) - ).toThrow( - "Add rows cannot be used with a named dataset for this element." - ) - }) - }) - - describe("arrowVegaLiteChart", () => { - const getVegaLiteChart = ( - datasets?: ArrowNamedDataSet[], - data?: Uint8Array - ): IArrowVegaLiteChart => ({ - datasets: datasets || [], - data: data ? { data } : null, - spec: JSON.stringify({ - mark: "circle", - encoding: { - x: { field: "a", type: "quantitative" }, - y: { field: "b", type: "quantitative" }, - size: { field: "c", type: "quantitative" }, - color: { field: "c", type: "quantitative" }, - }, - }), - useContainerWidth: true, - }) - - describe("addRows is called with a named dataset", () => { - test("element has one dataset -> append new rows to that dataset", () => { - const node = arrowVegaLiteChart( - getVegaLiteChart([MOCK_ANOTHER_NAMED_DATASET]) - ) - const newNode = node.arrowAddRows(MOCK_NAMED_DATASET, NO_SCRIPT_RUN_ID) - const element = newNode.vegaLiteChartElement - - const quiverData = element.datasets[0].data - expect(quiverData?.columnNames).toEqual([["", "c1", "c2"]]) - expect(quiverData?.dimensions.numDataRows).toEqual(4) - - expect(quiverData?.getCell(0, 0).content).toEqual("i1") - expect(quiverData?.getCell(0, 1).content).toEqual("foo") - expect(quiverData?.getCell(2, 0).content).toEqual("i1") - expect(quiverData?.getCell(2, 1).content).toEqual("foo") - }) - - test("element has a dataset with the given name -> append new rows to that dataset", () => { - const node = arrowVegaLiteChart( - getVegaLiteChart([MOCK_NAMED_DATASET, MOCK_ANOTHER_NAMED_DATASET]) - ) - const newNode = node.arrowAddRows(MOCK_NAMED_DATASET, NO_SCRIPT_RUN_ID) - const element = newNode.vegaLiteChartElement - - const quiverData = element.datasets[0].data - expect(quiverData?.columnNames).toEqual([["", "c1", "c2"]]) - expect(quiverData?.dimensions.numDataRows).toEqual(4) - - expect(quiverData?.getCell(0, 0).content).toEqual("i1") - expect(quiverData?.getCell(0, 1).content).toEqual("foo") - expect(quiverData?.getCell(2, 0).content).toEqual("i1") - expect(quiverData?.getCell(2, 1).content).toEqual("foo") - }) - - test("element doesn't have a matched dataset, but has data -> append new rows to data", () => { - const node = arrowVegaLiteChart(getVegaLiteChart(undefined, UNICODE)) - const newNode = node.arrowAddRows(MOCK_NAMED_DATASET, NO_SCRIPT_RUN_ID) - const element = newNode.vegaLiteChartElement - - const quiverData = element.data - expect(quiverData?.columnNames).toEqual([["", "c1", "c2"]]) - expect(quiverData?.dimensions.numDataRows).toEqual(4) - - expect(quiverData?.getCell(0, 0).content).toEqual("i1") - expect(quiverData?.getCell(0, 1).content).toEqual("foo") - expect(quiverData?.getCell(2, 0).content).toEqual("i1") - expect(quiverData?.getCell(2, 1).content).toEqual("foo") - }) - - test("element doesn't have a matched dataset or data -> use new rows as data", () => { - const node = arrowVegaLiteChart( - getVegaLiteChart([ - MOCK_ANOTHER_NAMED_DATASET, - MOCK_ANOTHER_NAMED_DATASET, - ]) - ) - const newNode = node.arrowAddRows(MOCK_NAMED_DATASET, NO_SCRIPT_RUN_ID) - const element = newNode.vegaLiteChartElement - - const quiverData = element.data - expect(quiverData?.columnNames).toEqual([["", "c1", "c2"]]) - expect(quiverData?.dimensions.numDataRows).toEqual(2) - - expect(quiverData?.getCell(0, 0).content).toEqual("i1") - expect(quiverData?.getCell(0, 1).content).toEqual("foo") - }) - - test("element doesn't have any datasets or data -> use new rows as data", () => { - const node = arrowVegaLiteChart(getVegaLiteChart()) - const newNode = node.arrowAddRows(MOCK_NAMED_DATASET, NO_SCRIPT_RUN_ID) - const element = newNode.vegaLiteChartElement - - const quiverData = element.data - expect(quiverData?.columnNames).toEqual([["", "c1", "c2"]]) - expect(quiverData?.dimensions.numDataRows).toEqual(2) - - expect(quiverData?.getCell(0, 0).content).toEqual("i1") - expect(quiverData?.getCell(0, 1).content).toEqual("foo") - }) - }) - - describe("addRows is called with an unnamed dataset", () => { - test("element has one dataset -> append new rows to that dataset", () => { - const node = arrowVegaLiteChart(getVegaLiteChart([MOCK_NAMED_DATASET])) - const newNode = node.arrowAddRows( - MOCK_UNNAMED_DATASET, - NO_SCRIPT_RUN_ID - ) - const element = newNode.vegaLiteChartElement - - const quiverData = element.datasets[0].data - expect(quiverData.columnNames).toEqual([["", "c1", "c2"]]) - expect(quiverData?.dimensions.numDataRows).toEqual(4) - - expect(quiverData.getCell(0, 0).content).toEqual("i1") - expect(quiverData.getCell(2, 0).content).toEqual("i1") - expect(quiverData.getCell(0, 1).content).toEqual("foo") - expect(quiverData.getCell(2, 1).content).toEqual("foo") - }) - - test("element has data -> append new rows to data", () => { - const node = arrowVegaLiteChart(getVegaLiteChart(undefined, UNICODE)) - const newNode = node.arrowAddRows( - MOCK_UNNAMED_DATASET, - NO_SCRIPT_RUN_ID - ) - const element = newNode.vegaLiteChartElement - - const quiverData = element.data - expect(quiverData?.columnNames).toEqual([["", "c1", "c2"]]) - expect(quiverData?.dimensions.numDataRows).toEqual(4) - - expect(quiverData?.getCell(0, 0).content).toEqual("i1") - expect(quiverData?.getCell(2, 0).content).toEqual("i1") - expect(quiverData?.getCell(0, 1).content).toEqual("foo") - expect(quiverData?.getCell(2, 1).content).toEqual("foo") - }) - - test("element doesn't have any datasets or data -> use new rows as data", () => { - const node = arrowVegaLiteChart(getVegaLiteChart()) - const newNode = node.arrowAddRows( - MOCK_UNNAMED_DATASET, - NO_SCRIPT_RUN_ID - ) - const element = newNode.vegaLiteChartElement - - const quiverData = element.data - expect(quiverData?.columnNames).toEqual([["", "c1", "c2"]]) - expect(quiverData?.dimensions.numDataRows).toEqual(2) - - expect(quiverData?.getCell(0, 0).content).toEqual("i1") - expect(quiverData?.getCell(0, 1).content).toEqual("foo") - }) - }) - }) - - it("throws an error for other element types", () => { - const node = text("foo") - expect(() => - node.arrowAddRows(MOCK_UNNAMED_DATASET, NO_SCRIPT_RUN_ID) - ).toThrow("elementType 'text' is not a valid arrowAddRows target!") - }) -}) - -describe("AppRoot.empty", () => { - let windowSpy: MockInstance - - beforeEach(() => { - windowSpy = vi.spyOn(window, "window", "get") - }) - - afterEach(() => { - windowSpy.mockRestore() - }) - - it("creates empty tree except for a skeleton", () => { - windowSpy.mockImplementation(() => ({ - location: { - search: "", - }, - })) - const empty = AppRoot.empty(FAKE_SCRIPT_HASH) - - expect(empty.main.children.length).toBe(1) - const child = empty.main.getIn([0]) as ElementNode - expect(child.element.skeleton).not.toBeNull() - - expect(empty.sidebar.isEmpty).toBe(true) - }) - - it("sets the main script hash and active script hash", () => { - windowSpy.mockImplementation(() => ({ - location: { - search: "", - }, - })) - const empty = AppRoot.empty(FAKE_SCRIPT_HASH) - - expect(empty.mainScriptHash).toBe(FAKE_SCRIPT_HASH) - expect(empty.main.activeScriptHash).toBe(FAKE_SCRIPT_HASH) - expect(empty.sidebar.activeScriptHash).toBe(FAKE_SCRIPT_HASH) - expect(empty.event.activeScriptHash).toBe(FAKE_SCRIPT_HASH) - expect(empty.bottom.activeScriptHash).toBe(FAKE_SCRIPT_HASH) - expect(empty.root.activeScriptHash).toBe(FAKE_SCRIPT_HASH) - }) - - it("creates empty tree with no loading screen if query param is set", () => { - windowSpy.mockImplementation(() => ({ - location: { - search: "?embed_options=hide_loading_screen", - }, - })) - - const empty = AppRoot.empty(FAKE_SCRIPT_HASH) - - expect(empty.main.isEmpty).toBe(true) - expect(empty.sidebar.isEmpty).toBe(true) - }) - - it("creates empty tree with v1 loading screen if query param is set", () => { - windowSpy.mockImplementation(() => ({ - location: { - search: "?embed_options=show_loading_screen_v1", - }, - })) - - const empty = AppRoot.empty(FAKE_SCRIPT_HASH) - - expect(empty.main.children.length).toBe(1) - const child = empty.main.getIn([0]) as ElementNode - expect(child.element.alert).toBeDefined() - - expect(empty.sidebar.isEmpty).toBe(true) - }) - - it("creates empty tree with v2 loading screen if query param is set", () => { - windowSpy.mockImplementation(() => ({ - location: { - search: "?embed_options=show_loading_screen_v2", - }, - })) - - const empty = AppRoot.empty(FAKE_SCRIPT_HASH) - - expect(empty.main.children.length).toBe(1) - const child = empty.main.getIn([0]) as ElementNode - expect(child.element.skeleton).not.toBeNull() - - expect(empty.sidebar.isEmpty).toBe(true) - }) - - it("creates empty tree with no loading screen if query param is v1 and it's not first load", () => { - windowSpy.mockImplementation(() => ({ - location: { - search: "?embed_options=show_loading_screen_v1", - }, - })) - - const empty = AppRoot.empty(FAKE_SCRIPT_HASH, false) - - expect(empty.main.isEmpty).toBe(true) - expect(empty.sidebar.isEmpty).toBe(true) - }) - - it("passes logo to new Root if empty is called with logo", () => { - windowSpy.mockImplementation(() => ({ - location: { - search: "", - }, - })) - const logo = LogoProto.create({ - image: - "https://global.discourse-cdn.com/business7/uploads/streamlit/original/2X/8/8cb5b6c0e1fe4e4ebfd30b769204c0d30c332fec.png", - }) - - // Replicate .empty call on initial render - const empty = AppRoot.empty("", true) - expect(empty.logo).toBeNull() - - // Replicate .empty call in AppNav's clearPageElements for MPA V1 - const empty2 = AppRoot.empty(FAKE_SCRIPT_HASH, false, undefined, logo) - expect(empty2.logo).not.toBeNull() - }) -}) - -describe("AppRoot.filterMainScriptElements", () => { - it("does not clear nodes associated with main script hash", () => { - // Add a new element and clear stale nodes - const delta = makeProto(DeltaProto, { - newElement: { text: { body: "newElement!" } }, - }) - const newRoot = ROOT.applyDelta( - "new_session_id", - delta, - forwardMsgMetadata([0, 1, 1]) - ).filterMainScriptElements(FAKE_SCRIPT_HASH) - - // We should now only have a single element, inside a single block - expect(newRoot.main.getIn([1, 1])).toBeTextNode("newElement!") - expect(newRoot.getElements().size).toBe(3) - }) - - it("clears nodes not associated with main script hash", () => { - // Add a new element and clear stale nodes - const delta = makeProto(DeltaProto, { - newElement: { text: { body: "newElement!" } }, - }) - const newRoot = ROOT.applyDelta( - "new_session_id", - delta, - forwardMsgMetadata([0, 1, 1], "DIFFERENT_HASH") - ).filterMainScriptElements(FAKE_SCRIPT_HASH) - - // We should now only have a single element, inside a single block - expect(newRoot.main.getIn([1, 1])).toBeUndefined() - expect(newRoot.getElements().size).toBe(2) - }) -}) - -describe("AppRoot.applyDelta", () => { - it("handles 'newElement' deltas", () => { - const delta = makeProto(DeltaProto, { - newElement: { text: { body: "newElement!" } }, - }) - const newRoot = ROOT.applyDelta( - "new_session_id", - delta, - forwardMsgMetadata([0, 1, 1]) - ) - - const newNode = newRoot.main.getIn([1, 1]) as ElementNode - expect(newNode).toBeTextNode("newElement!") - - // Check that our new scriptRunId has been set only on the touched nodes - expect(newRoot.main.scriptRunId).toBe("new_session_id") - expect(newRoot.main.fragmentId).toBe(undefined) - expect(newRoot.main.deltaMsgReceivedAt).toBe(undefined) - expect(newRoot.main.getIn([0])?.scriptRunId).toBe(NO_SCRIPT_RUN_ID) - expect(newRoot.main.getIn([1])?.scriptRunId).toBe("new_session_id") - expect(newRoot.main.getIn([1, 0])?.scriptRunId).toBe(NO_SCRIPT_RUN_ID) - expect(newRoot.main.getIn([1, 1])?.scriptRunId).toBe("new_session_id") - expect(newNode.activeScriptHash).toBe(FAKE_SCRIPT_HASH) - expect(newRoot.sidebar.scriptRunId).toBe(NO_SCRIPT_RUN_ID) - }) - - it("handles 'addBlock' deltas", () => { - const delta = makeProto(DeltaProto, { addBlock: {} }) - const newRoot = ROOT.applyDelta( - "new_session_id", - delta, - forwardMsgMetadata([0, 1, 1]) - ) - - const newNode = newRoot.main.getIn([1, 1]) as BlockNode - expect(newNode).toBeDefined() - - // Check that our new scriptRunId has been set only on the touched nodes - expect(newRoot.main.scriptRunId).toBe("new_session_id") - expect(newRoot.main.fragmentId).toBe(undefined) - expect(newRoot.main.deltaMsgReceivedAt).toBe(undefined) - expect(newRoot.main.getIn([0])?.scriptRunId).toBe(NO_SCRIPT_RUN_ID) - expect(newRoot.main.getIn([1])?.scriptRunId).toBe("new_session_id") - expect(newRoot.main.getIn([1, 0])?.scriptRunId).toBe(NO_SCRIPT_RUN_ID) - expect(newRoot.main.getIn([1, 1])?.scriptRunId).toBe("new_session_id") - expect(newNode.activeScriptHash).toBe(FAKE_SCRIPT_HASH) - expect(newRoot.sidebar.scriptRunId).toBe(NO_SCRIPT_RUN_ID) - }) - - it("removes a block's children if the block type changes for the same delta path", () => { - const newRoot = ROOT.applyDelta( - "script_run_id", - makeProto(DeltaProto, { - addBlock: { - expandable: { - expanded: true, - label: "label", - icon: "", - }, - }, - }), - forwardMsgMetadata([0, 1, 1]) - ).applyDelta( - "script_run_id", - makeProto(DeltaProto, { - newElement: { text: { body: "newElement!" } }, - }), - forwardMsgMetadata([0, 1, 1, 0]) - ) - - const newNode = newRoot.main.getIn([1, 1]) as BlockNode - expect(newNode).toBeDefined() - expect(newNode.deltaBlock.type).toBe("expandable") - expect(newNode.children.length).toBe(1) - - const newRoot2 = newRoot.applyDelta( - "new_script_run_id", - makeProto(DeltaProto, { - addBlock: { - tabContainer: {}, - }, - }), - forwardMsgMetadata([0, 1, 1]) - ) - - const replacedBlock = newRoot2.main.getIn([1, 1]) as BlockNode - expect(replacedBlock).toBeDefined() - expect(replacedBlock.deltaBlock.type).toBe("tabContainer") - expect(replacedBlock.children.length).toBe(0) - }) - - it("will not remove a block's children if the block type is the same for the same delta path", () => { - const newRoot = ROOT.applyDelta( - "script_run_id", - makeProto(DeltaProto, { - addBlock: { - expandable: { - expanded: true, - label: "label", - icon: "", - }, - }, - }), - forwardMsgMetadata([0, 1, 1]) - ).applyDelta( - "script_run_id", - makeProto(DeltaProto, { - newElement: { text: { body: "newElement!" } }, - }), - forwardMsgMetadata([0, 1, 1, 0]) - ) - - const newNode = newRoot.main.getIn([1, 1]) as BlockNode - expect(newNode).toBeDefined() - expect(newNode.deltaBlock.type).toBe("expandable") - expect(newNode.children.length).toBe(1) - - const newRoot2 = newRoot.applyDelta( - "new_script_run_id", - makeProto(DeltaProto, { - addBlock: { - expandable: { - expanded: true, - label: "other label", - icon: "", - }, - }, - }), - forwardMsgMetadata([0, 1, 1]) - ) - - const replacedBlock = newRoot2.main.getIn([1, 1]) as BlockNode - expect(replacedBlock).toBeDefined() - expect(replacedBlock.deltaBlock.type).toBe("expandable") - expect(replacedBlock.children.length).toBe(1) - }) - - it("specifies active script hash on 'newElement' deltas", () => { - const delta = makeProto(DeltaProto, { - newElement: { text: { body: "newElement!" } }, - }) - const NEW_FAKE_SCRIPT_HASH = "new_fake_script_hash" - const newRoot = ROOT.applyDelta( - "new_session_id", - delta, - forwardMsgMetadata([0, 1, 1], NEW_FAKE_SCRIPT_HASH) - ) - - const newNode = newRoot.main.getIn([1, 1]) as ElementNode - expect(newNode).toBeDefined() - - // Check that our new other nodes are not affected by the new script hash - expect(newRoot.main.getIn([1, 0])?.activeScriptHash).toBe(FAKE_SCRIPT_HASH) - expect(newNode.activeScriptHash).toBe(NEW_FAKE_SCRIPT_HASH) - }) - - it("specifies active script hash on 'addBlock' deltas", () => { - const delta = makeProto(DeltaProto, { addBlock: {} }) - const NEW_FAKE_SCRIPT_HASH = "new_fake_script_hash" - const newRoot = ROOT.applyDelta( - "new_session_id", - delta, - forwardMsgMetadata([0, 1, 1], NEW_FAKE_SCRIPT_HASH) - ) - - const newNode = newRoot.main.getIn([1, 1]) as BlockNode - expect(newNode).toBeDefined() - - // Check that our new scriptRunId has been set only on the touched nodes - expect(newRoot.main.getIn([1, 0])?.activeScriptHash).toBe(FAKE_SCRIPT_HASH) - expect(newNode.activeScriptHash).toBe(NEW_FAKE_SCRIPT_HASH) - }) - - it("can set fragmentId in 'newElement' deltas", () => { - const delta = makeProto(DeltaProto, { - newElement: { text: { body: "newElement!" } }, - fragmentId: "myFragmentId", - }) - const newRoot = ROOT.applyDelta( - "new_session_id", - delta, - forwardMsgMetadata([0, 1, 1]) - ) - - const newNode = newRoot.main.getIn([1, 1]) as ElementNode - expect(newNode.fragmentId).toBe("myFragmentId") - }) - - it("can set fragmentId in 'addBlock' deltas", () => { - const delta = makeProto(DeltaProto, { - addBlock: {}, - fragmentId: "myFragmentId", - }) - const newRoot = ROOT.applyDelta( - "new_session_id", - delta, - forwardMsgMetadata([0, 1, 1]) - ) - - const newNode = newRoot.main.getIn([1, 1]) as BlockNode - expect(newNode.fragmentId).toBe("myFragmentId") - }) - - it("timestamp is set on BlockNode as message id", () => { - const timestamp = new Date(Date.UTC(2017, 1, 14)).valueOf() - Date.now = vi.fn(() => timestamp) - const delta = makeProto(DeltaProto, { - addBlock: {}, - }) - const newRoot = ROOT.applyDelta( - "new_session_id", - delta, - forwardMsgMetadata([0, 1, 1]) - ) - - const newNode = newRoot.main.getIn([1, 1]) as BlockNode - expect(newNode.deltaMsgReceivedAt).toBe(timestamp) - }) -}) - -describe("AppRoot.clearStaleNodes", () => { - it("clears stale nodes", () => { - // Add a new element and clear stale nodes - const delta = makeProto(DeltaProto, { - newElement: { text: { body: "newElement!" } }, - }) - const newRoot = ROOT.applyDelta( - "new_session_id", - delta, - forwardMsgMetadata([0, 1, 1]) - ).clearStaleNodes("new_session_id", []) - - // We should now only have a single element, inside a single block - expect(newRoot.main.getIn([0, 0])).toBeTextNode("newElement!") - expect(newRoot.getElements().size).toBe(1) - }) - - it("clears a stale logo", () => { - const logo = LogoProto.create({ - image: - "https://global.discourse-cdn.com/business7/uploads/streamlit/original/2X/8/8cb5b6c0e1fe4e4ebfd30b769204c0d30c332fec.png", - }) - const newRoot = ROOT.appRootWithLogo(logo, { - activeScriptHash: "hash", - scriptRunId: "script_run_id", - }) - expect(newRoot.logo).not.toBeNull() - - const newNewRoot = newRoot.clearStaleNodes("new_script_run_id", []) - expect(newNewRoot.logo).toBeNull() - }) - - it("does not clear logo on fragment run", () => { - const logo = LogoProto.create({ - image: - "https://global.discourse-cdn.com/business7/uploads/streamlit/original/2X/8/8cb5b6c0e1fe4e4ebfd30b769204c0d30c332fec.png", - }) - const newRoot = ROOT.appRootWithLogo(logo, { - activeScriptHash: "hash", - scriptRunId: "script_run_id", - }) - expect(newRoot.logo).not.toBeNull() - - const newNewRoot = newRoot.clearStaleNodes("new_script_run_id", [ - "my_fragment_id", - ]) - expect(newNewRoot.logo).not.toBeNull() - }) - - it("handles currentFragmentId correctly", () => { - const tabContainerProto = makeProto(DeltaProto, { - addBlock: { tabContainer: {}, allowEmpty: false }, - fragmentId: "my_fragment_id", - }) - const tab1 = makeProto(DeltaProto, { - addBlock: { tab: { label: "tab1" }, allowEmpty: true }, - fragmentId: "my_fragment_id", - }) - const tab2 = makeProto(DeltaProto, { - addBlock: { tab: { label: "tab2" }, allowEmpty: true }, - fragmentId: "my_fragment_id", - }) - - // const BLOCK = block([text("1"), block([text("2")])]) - const root = AppRoot.empty(FAKE_SCRIPT_HASH) - // Block not corresponding to my_fragment_id. Should be preserved. - .applyDelta( - "old_session_id", - makeProto(DeltaProto, { addBlock: { allowEmpty: true } }), - forwardMsgMetadata([0, 0]) - ) - // Element in block unrelated to my_fragment_id. Should be preserved. - .applyDelta( - "old_session_id", - makeProto(DeltaProto, { - newElement: { text: { body: "oldElement!" } }, - }), - forwardMsgMetadata([0, 0, 0]) - ) - // Another element in block unrelated to my_fragment_id. Should be preserved. - .applyDelta( - "old_session_id", - makeProto(DeltaProto, { - newElement: { text: { body: "oldElement2!" } }, - fragmentId: "other_fragment_id", - }), - forwardMsgMetadata([0, 0, 1]) - ) - // Old element related to my_fragment_id but in an unrelated block. Should be preserved. - .applyDelta( - "old_session_id", - makeProto(DeltaProto, { - newElement: { text: { body: "oldElement4!" } }, - fragmentId: "my_fragment_id", - }), - forwardMsgMetadata([0, 0, 2]) - ) - // Block corresponding to my_fragment_id - .applyDelta( - "new_session_id", - makeProto(DeltaProto, { - addBlock: { allowEmpty: false }, - fragmentId: "my_fragment_id", - }), - forwardMsgMetadata([0, 1]) - ) - // Old element related to my_fragment_id. Should be pruned. - .applyDelta( - "old_session_id", - makeProto(DeltaProto, { - newElement: { text: { body: "oldElement3!" } }, - fragmentId: "my_fragment_id", - }), - forwardMsgMetadata([0, 1, 0]) - ) - // New element related to my_fragment_id. Should be preserved. - .applyDelta( - "new_session_id", - makeProto(DeltaProto, { - newElement: { text: { body: "newElement!" } }, - fragmentId: "my_fragment_id", - }), - forwardMsgMetadata([0, 1, 1]) - ) - // New element container related to my_fragment_id, having children which will be handled individually - // Create a tab container with two tabs in the old session; then send new delta with the container and - // only one tab. The second tab with the old_session_id should be pruned. - .applyDelta( - "old_session_id", - tabContainerProto, - forwardMsgMetadata([0, 2]) - ) - .applyDelta("old_session_id", tab1, forwardMsgMetadata([0, 2, 0])) - .applyDelta("old_session_id", tab2, forwardMsgMetadata([0, 2, 1])) - .applyDelta( - "new_session_id", - tabContainerProto, - forwardMsgMetadata([0, 2]) - ) - .applyDelta("new_session_id", tab1, forwardMsgMetadata([0, 2, 0])) - - const pruned = root.clearStaleNodes("new_session_id", ["my_fragment_id"]) - - expect(pruned.main.getIn([0])).toBeInstanceOf(BlockNode) - expect((pruned.main.getIn([0]) as BlockNode).children).toHaveLength(3) - expect(pruned.main.getIn([0, 0])).toBeTextNode("oldElement!") - expect(pruned.main.getIn([0, 1])).toBeTextNode("oldElement2!") - expect(pruned.main.getIn([0, 2])).toBeTextNode("oldElement4!") - - expect(pruned.main.getIn([1])).toBeInstanceOf(BlockNode) - expect((pruned.main.getIn([1]) as BlockNode).children).toHaveLength(1) - expect(pruned.main.getIn([1, 0])).toBeTextNode("newElement!") - - expect(pruned.main.getIn([2])).toBeInstanceOf(BlockNode) - expect((pruned.main.getIn([2]) as BlockNode).children).toHaveLength(1) - expect( - (pruned.main.getIn([2, 0]) as BlockNode).deltaBlock.tab?.label - ).toContain("tab1") - }) - - it("clear childNodes of a block node in fragment run", () => { - // Add a new element and clear stale nodes - const delta = makeProto(DeltaProto, { - newElement: { text: { body: "newElement!" } }, - fragmentId: "my_fragment_id", - }) - const newRoot = AppRoot.empty(FAKE_SCRIPT_HASH) - // Block corresponding to my_fragment_id - .applyDelta( - "new_session_id", - makeProto(DeltaProto, { - addBlock: { vertical: {}, allowEmpty: false }, - fragmentId: "my_fragment_id", - }), - forwardMsgMetadata([0, 0]) - ) - .applyDelta("new_session_id", delta, forwardMsgMetadata([0, 0, 0])) - // Block with child where scriptRunId is different - .applyDelta( - "new_session_id", - makeProto(DeltaProto, { - addBlock: { vertical: {}, allowEmpty: false }, - fragmentId: "my_fragment_id", - }), - forwardMsgMetadata([0, 1]) - ) - .applyDelta("new_session_id", delta, forwardMsgMetadata([0, 1, 0])) - .applyDelta("new_session_id", delta, forwardMsgMetadata([0, 1, 1])) - // this child is a nested fragment_id from an old run and should be pruned - .applyDelta( - "old_session_id", - makeProto(DeltaProto, { - newElement: { text: { body: "oldElement!" } }, - fragmentId: "my_nested_fragment_id", - }), - forwardMsgMetadata([0, 1, 2]) - ) - // this child is a nested fragment_id from the same run and should be preserved - .applyDelta( - "new_session_id", - makeProto(DeltaProto, { - newElement: { text: { body: "newElement!" } }, - fragmentId: "my_nested_fragment_id", - }), - forwardMsgMetadata([0, 1, 3]) - ) - - expect((newRoot.main.getIn([1]) as BlockNode).children).toHaveLength(4) - - const pruned = newRoot.clearStaleNodes("new_session_id", [ - "my_fragment_id", - ]) - - expect(pruned.main.getIn([0])).toBeInstanceOf(BlockNode) - expect((pruned.main.getIn([0]) as BlockNode).children).toHaveLength(1) - expect(pruned.main.getIn([1])).toBeInstanceOf(BlockNode) - // the stale nested fragment child should have been pruned - expect((pruned.main.getIn([1]) as BlockNode).children).toHaveLength(3) - }) -}) - -describe("AppRoot.getElements", () => { - it("returns all elements", () => { - // We have elements at main.[0] and main.[1, 0] - expect(ROOT.getElements()).toEqual( - new Set([ - (ROOT.main.getIn([0]) as ElementNode).element, - (ROOT.main.getIn([1, 0]) as ElementNode).element, - ]) - ) - }) -}) - -/** Create a `Text` element node with the given properties. */ -function text(textArg: string, scriptRunId = NO_SCRIPT_RUN_ID): ElementNode { - const element = makeProto(Element, { text: { body: textArg } }) - return new ElementNode( - element, - ForwardMsgMetadata.create(), - scriptRunId, - FAKE_SCRIPT_HASH - ) -} - -/** Create a BlockNode with the given properties. */ -function block( - children: AppNode[] = [], - scriptRunId = NO_SCRIPT_RUN_ID -): BlockNode { - return new BlockNode( - FAKE_SCRIPT_HASH, - children, - makeProto(BlockProto, {}), - scriptRunId - ) -} - -/** Create an arrowTable element node with the given properties. */ -function arrowTable(scriptRunId = NO_SCRIPT_RUN_ID): ElementNode { - const element = makeProto(Element, { arrowTable: { data: UNICODE } }) - return new ElementNode( - element, - ForwardMsgMetadata.create(), - scriptRunId, - FAKE_SCRIPT_HASH - ) -} - -/** Create an arrowDataFrame element node with the given properties. */ -function arrowDataFrame(scriptRunId = NO_SCRIPT_RUN_ID): ElementNode { - const element = makeProto(Element, { arrowDataFrame: { data: UNICODE } }) - return new ElementNode( - element, - ForwardMsgMetadata.create(), - scriptRunId, - FAKE_SCRIPT_HASH - ) -} - -/** Create an arrowVegaLiteChart element node with the given properties. */ -function arrowVegaLiteChart( - data: IArrowVegaLiteChart, - scriptRunId = NO_SCRIPT_RUN_ID -): ElementNode { - const element = makeProto(Element, { arrowVegaLiteChart: data }) - return new ElementNode( - element, - ForwardMsgMetadata.create(), - scriptRunId, - FAKE_SCRIPT_HASH - ) -} - -/** Create a ForwardMsgMetadata with the given container and path */ -function forwardMsgMetadata( - deltaPath: number[], - activeScriptHash = FAKE_SCRIPT_HASH -): ForwardMsgMetadata { - expect(deltaPath.length).toBeGreaterThanOrEqual(2) - return makeProto(ForwardMsgMetadata, { deltaPath, activeScriptHash }) -} - -/** - * Make a "fully concrete" instance of a protobuf message. - * This function constructs a message and then encodes and decodes it as - * if it had arrived on the wire. This ensures that that it has all its - * 'oneOfs' and 'defaults' set. - */ -function makeProto( - MessageType: { - new (props: Props): Type - encode: (message: Type, writer: Writer) => Writer - decode: (bytes: Uint8Array) => Type - }, - properties: Props -): Type { - const message = new MessageType(properties) - const bytes = MessageType.encode(message, Writer.create()).finish() - return MessageType.decode(bytes) -} - -// Custom Jest matchers for dealing with AppNodes -declare global { - // eslint-disable-next-line @typescript-eslint/no-namespace - namespace vi { - interface Matchers { - toBeTextNode(text: string): R - } - } -} - -interface CustomMatchers { - toBeTextNode(text: string): R -} - -declare module "vitest" { - // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-empty-object-type -- TODO: Replace 'any' with a more specific type. - interface Assertion extends CustomMatchers {} - // eslint-disable-next-line @typescript-eslint/no-empty-object-type - interface AsymmetricMatchersContaining extends CustomMatchers {} -} - -expect.extend({ - toBeTextNode(received, textArg) { - const elementNode = received as ElementNode - if (isNullOrUndefined(elementNode)) { - return { - message: () => `expected ${received} to be an instance of ElementNode`, - pass: false, - } - } - - const { type } = elementNode.element - if (type !== "text") { - return { - message: () => - `expected ${received}.element.type to be 'text', but it was ${type}`, - pass: false, - } - } - - const textBody = elementNode.element.text?.body - return { - message: () => - `expected ${received}.element.text.body to be "${textArg}", but it was "${textBody}"`, - pass: textBody === textArg, - } - }, -}) diff --git a/frontend/lib/src/AppNode.ts b/frontend/lib/src/AppNode.ts index 42830b9e8c1..ab6ca8f9183 100644 --- a/frontend/lib/src/AppNode.ts +++ b/frontend/lib/src/AppNode.ts @@ -14,928 +14,7 @@ * limitations under the License. */ -import { produce } from "immer" - -import { - ArrowNamedDataSet, - Arrow as ArrowProto, - ArrowVegaLiteChart as ArrowVegaLiteChartProto, - Block as BlockProto, - Delta, - Element, - ForwardMsgMetadata, - IArrow, - IArrowNamedDataSet, - Logo, -} from "@streamlit/protobuf" - -import { - getLoadingScreenType, - isNullOrUndefined, - LoadingScreenType, - makeAppSkeletonElement, - makeElementWithErrorText, - makeElementWithInfoText, - notUndefined, -} from "~lib/util/utils" - -import { - VegaLiteChartElement, - WrappedNamedDataset, -} from "./components/elements/ArrowVegaLiteChart" -import { Quiver } from "./dataframes/Quiver" -import { ensureError } from "./util/ErrorHandling" - -const NO_SCRIPT_RUN_ID = "NO_SCRIPT_RUN_ID" - -interface LogoMetadata { - // Associated scriptHash that created the logo - activeScriptHash: string - - // Associated scriptRunId that created the logo - scriptRunId: string -} -interface AppLogo extends LogoMetadata { - logo: Logo -} - -/** - * An immutable node of the "App Data Tree". - * - * Trees are composed of `ElementNode` leaves, which contain data about - * a single visual element, and `BlockNode` branches, which determine the - * layout of a group of children nodes. - * - * A simple tree might look like this: - * - * AppRoot - * ├── BlockNode ("main") - * │ ├── ElementNode (text: "Ahoy, Streamlit!") - * │ └── ElementNode (button: "Don't Push This") - * └── BlockNode ("sidebar") - * └── ElementNode (checkbox: "Batten The Hatches") - * - * To build this tree, the frontend receives `Delta` messages from Python, - * each of which corresponds to a tree mutation ("add an element", - * "add a block", "add rows to an existing element"). The frontend builds the - * tree bit by bit in response to these `Delta`s. - * - * To render the app, the `AppView` class walks this tree and outputs - * a corresponding DOM structure, using React, that's essentially a mapping - * of `AppElement` -> `ReactNode`. This rendering happens "live" - that is, - * the app is re-rendered each time a new `Delta` is received. - * - * Because the app gets re-rendered frequently, updates need to be fast. - * Our React components - the building blocks of the app - are "pure" - * (see https://reactjs.org/docs/react-api.html#reactpurecomponent), which - * means that React uses shallow comparison to determine which ReactNodes to - * update. - * - * Thus, each node in our tree is _immutable_ - any change to a `AppNode` - * actually results in a *new* `AppNode` instance. This occurs recursively, - * so inserting a new `ElementNode` into the tree will also result in new - * `BlockNode`s for each of that Element's ancestors, all the way up to the - * root node. Then, when React re-renders the app, it will re-traverse the new - * nodes that have been created, and rebuild just the bits of the app that - * have changed. - */ -export interface AppNode { - /** - * The ID of the script run this node was generated in. When a script finishes - * running, the app prunes all stale nodes. - */ - readonly scriptRunId: string - - /** - * The ID of the fragment that sent the Delta creating this AppNode. If this - * AppNode was not created by a fragment, this field is falsy. - */ - readonly fragmentId?: string - - /** - * The hash of the script that created this node. - */ - readonly activeScriptHash?: string - - // A timestamp indicating based on which delta message the node was created. - // If the node was created without a delta message, this field is undefined. - // This helps us to update React components based on a new backend message even though other - // props have not changed; this can happen for UI-only interactions such as dismissing a dialog. - readonly deltaMsgReceivedAt?: number - - /** - * Return the AppNode for the given index path, or undefined if the path - * is invalid. - */ - getIn(path: number[]): AppNode | undefined - - /** - * Return a copy of this node with a new element set at the given index - * path. Throws an error if the path is invalid. - */ - setIn(path: number[], node: AppNode, scriptRunId: string): AppNode - - /** - * Recursively remove children nodes whose activeScriptHash is no longer - * associated with the mainScriptHash. - */ - filterMainScriptElements(mainScriptHash: string): AppNode | undefined - - /** - * Recursively remove children nodes whose scriptRunId is no longer current. - * If this node should no longer exist, return undefined. - */ - clearStaleNodes( - currentScriptRunId: string, - fragmentIdsThisRun?: Array, - fragmentIdOfBlock?: string - ): AppNode | undefined - - /** - * Return a Set of all the Elements contained in the tree. - * If an existing Set is passed in, that Set will be mutated and returned. - * Otherwise, a new Set will be created and will be returned. - */ - getElements(elementSet?: Set): Set -} - -/** - * A leaf AppNode. Contains a single element to render. - */ -export class ElementNode implements AppNode { - public readonly element: Element - - public readonly metadata: ForwardMsgMetadata - - public readonly scriptRunId: string - - public readonly fragmentId?: string - - private lazyQuiverElement?: Quiver - - private lazyVegaLiteChartElement?: VegaLiteChartElement - - // The hash of the script that created this element. - public readonly activeScriptHash: string - - /** Create a new ElementNode. */ - public constructor( - element: Element, - metadata: ForwardMsgMetadata, - scriptRunId: string, - activeScriptHash: string, - fragmentId?: string - ) { - this.element = element - this.metadata = metadata - this.scriptRunId = scriptRunId - this.activeScriptHash = activeScriptHash - this.fragmentId = fragmentId - } - - public get quiverElement(): Quiver { - if (this.lazyQuiverElement !== undefined) { - return this.lazyQuiverElement - } - - if ( - this.element.type !== "arrowTable" && - this.element.type !== "arrowDataFrame" - ) { - throw new Error( - `elementType '${this.element.type}' is not a valid Quiver element!` - ) - } - - const toReturn = new Quiver(this.element[this.element.type] as ArrowProto) - // TODO (lukasmasuch): Delete element from proto object? - this.lazyQuiverElement = toReturn - return toReturn - } - - public get vegaLiteChartElement(): VegaLiteChartElement { - if (this.lazyVegaLiteChartElement !== undefined) { - return this.lazyVegaLiteChartElement - } - - if (this.element.type !== "arrowVegaLiteChart") { - throw new Error( - `elementType '${this.element.type}' is not a valid VegaLiteChartElement!` - ) - } - - const proto = this.element.arrowVegaLiteChart as ArrowVegaLiteChartProto - const modifiedData = proto.data ? new Quiver(proto.data) : null - const modifiedDatasets = - proto.datasets.length > 0 ? wrapDatasets(proto.datasets) : [] - - const toReturn = { - data: modifiedData, - spec: proto.spec, - datasets: modifiedDatasets, - useContainerWidth: proto.useContainerWidth, - vegaLiteTheme: proto.theme, - id: proto.id, - selectionMode: proto.selectionMode, - formId: proto.formId, - } - - this.lazyVegaLiteChartElement = toReturn - return toReturn - } - - public getIn(): AppNode | undefined { - return undefined - } - - public setIn(): AppNode { - throw new Error("'setIn' cannot be called on an ElementNode") - } - - public filterMainScriptElements( - mainScriptHash: string - ): AppNode | undefined { - if (this.activeScriptHash !== mainScriptHash) { - return undefined - } - - return this - } - - public clearStaleNodes( - currentScriptRunId: string, - fragmentIdsThisRun?: Array, - fragmentIdOfBlock?: string - ): ElementNode | undefined { - if (fragmentIdsThisRun?.length) { - // If we're currently running a fragment, nodes unrelated to the fragment - // shouldn't be cleared. This can happen when, - // 1. This element doesn't correspond to a fragment at all. - // 2. This element is a fragment but is in no path that was modified. - // 3. This element belongs to a path that was modified, but it was modified in the same run. - if ( - !this.fragmentId || - !fragmentIdOfBlock || - this.scriptRunId === currentScriptRunId - ) { - return this - } - } - return this.scriptRunId === currentScriptRunId ? this : undefined - } - - public getElements(elements?: Set): Set { - if (isNullOrUndefined(elements)) { - elements = new Set() - } - elements.add(this.element) - return elements - } - - public arrowAddRows( - namedDataSet: ArrowNamedDataSet, - scriptRunId: string - ): ElementNode { - const elementType = this.element.type - const newNode = new ElementNode( - this.element, - this.metadata, - scriptRunId, - this.activeScriptHash, - this.fragmentId - ) - - switch (elementType) { - case "arrowTable": - case "arrowDataFrame": { - newNode.lazyQuiverElement = ElementNode.quiverAddRowsHelper( - this.quiverElement, - namedDataSet - ) - break - } - case "arrowVegaLiteChart": { - newNode.lazyVegaLiteChartElement = - ElementNode.vegaLiteChartAddRowsHelper( - this.vegaLiteChartElement, - namedDataSet - ) - break - } - default: { - // This should never happen! - throw new Error( - `elementType '${this.element.type}' is not a valid arrowAddRows target!` - ) - } - } - - return newNode - } - - private static quiverAddRowsHelper( - element: Quiver, - namedDataSet: ArrowNamedDataSet - ): Quiver { - if (namedDataSet.hasName) { - throw new Error( - "Add rows cannot be used with a named dataset for this element." - ) - } - - const newQuiver = new Quiver(namedDataSet.data as IArrow) - return element.addRows(newQuiver) - } - - private static vegaLiteChartAddRowsHelper( - element: VegaLiteChartElement, - namedDataSet: ArrowNamedDataSet - ): VegaLiteChartElement { - const newDataSetName = namedDataSet.hasName ? namedDataSet.name : null - const newDataSetQuiver = new Quiver(namedDataSet.data as IArrow) - - return produce(element, (draft: VegaLiteChartElement) => { - const existingDataSet = getNamedDataSet(draft.datasets, newDataSetName) - if (existingDataSet) { - existingDataSet.data = existingDataSet.data.addRows(newDataSetQuiver) - } else { - draft.data = draft.data - ? draft.data.addRows(newDataSetQuiver) - : newDataSetQuiver - } - }) - } -} - -/** - * If there is only one NamedDataSet, return it. - * If there is a NamedDataset that matches the given name, return it. - * Otherwise, return `undefined`. - */ -function getNamedDataSet( - namedDataSets: WrappedNamedDataset[], - name: string | null -): WrappedNamedDataset | undefined { - if (namedDataSets.length === 1) { - return namedDataSets[0] - } - - return namedDataSets.find( - (dataset: WrappedNamedDataset) => dataset.hasName && dataset.name === name - ) -} - -/** - * A container AppNode that holds children. - */ -export class BlockNode implements AppNode { - public readonly children: AppNode[] - - public readonly deltaBlock: BlockProto - - public readonly scriptRunId: string - - public readonly fragmentId?: string - - public readonly deltaMsgReceivedAt?: number - - // The hash of the script that created this block. - public readonly activeScriptHash: string - - public constructor( - activeScriptHash: string, - children?: AppNode[], - deltaBlock?: BlockProto, - scriptRunId?: string, - fragmentId?: string, - deltaMsgReceivedAt?: number - ) { - this.activeScriptHash = activeScriptHash - this.children = children ?? [] - this.deltaBlock = deltaBlock ?? new BlockProto({}) - this.scriptRunId = scriptRunId ?? NO_SCRIPT_RUN_ID - this.fragmentId = fragmentId - this.deltaMsgReceivedAt = deltaMsgReceivedAt - } - - /** True if this Block has no children. */ - public get isEmpty(): boolean { - return this.children.length === 0 - } - - public getIn(path: number[]): AppNode | undefined { - if (path.length === 0) { - return undefined - } - - const childIndex = path[0] - if (childIndex < 0 || childIndex >= this.children.length) { - return undefined - } - - if (path.length === 1) { - return this.children[childIndex] - } - - return this.children[childIndex].getIn(path.slice(1)) - } - - public setIn(path: number[], node: AppNode, scriptRunId: string): BlockNode { - if (path.length === 0) { - throw new Error(`empty path!`) - } - - const childIndex = path[0] - if (childIndex < 0 || childIndex > this.children.length) { - throw new Error( - `Bad 'setIn' index ${childIndex} (should be between [0, ${this.children.length}])` - ) - } - - const newChildren = this.children.slice() - if (path.length === 1) { - // Base case - newChildren[childIndex] = node - } else { - // Pop the current element off our path, and recurse into our children - newChildren[childIndex] = newChildren[childIndex].setIn( - path.slice(1), - node, - scriptRunId - ) - } - - return new BlockNode( - this.activeScriptHash, - newChildren, - this.deltaBlock, - scriptRunId, - this.fragmentId, - this.deltaMsgReceivedAt - ) - } - - filterMainScriptElements(mainScriptHash: string): AppNode | undefined { - if (this.activeScriptHash !== mainScriptHash) { - return undefined - } - - // Recursively clear our children. - const newChildren = this.children - .map(child => child.filterMainScriptElements(mainScriptHash)) - .filter(notUndefined) - - return new BlockNode( - this.activeScriptHash, - newChildren, - this.deltaBlock, - this.scriptRunId, - this.fragmentId, - this.deltaMsgReceivedAt - ) - } - - public clearStaleNodes( - currentScriptRunId: string, - fragmentIdsThisRun?: Array, - fragmentIdOfBlock?: string - ): BlockNode | undefined { - if (!fragmentIdsThisRun?.length) { - // If we're not currently running a fragment, then we can remove any blocks - // that don't correspond to currentScriptRunId. - if (this.scriptRunId !== currentScriptRunId) { - return undefined - } - } else { - // Otherwise, we are currently running a fragment, and our behavior - // depends on the fragmentId of this BlockNode. - - // The parent block was modified but this element wasn't, so it's stale. - if (fragmentIdOfBlock && this.scriptRunId !== currentScriptRunId) { - return undefined - } - - // This block is modified by the current run, so we indicate this to our children in case - // they were not modified by the current run, which means they are stale. - if ( - this.fragmentId && - fragmentIdsThisRun.includes(this.fragmentId) && - this.scriptRunId === currentScriptRunId - ) { - fragmentIdOfBlock = this.fragmentId - } - } - - // Recursively clear our children. - const newChildren = this.children - .map(child => { - return child.clearStaleNodes( - currentScriptRunId, - fragmentIdsThisRun, - fragmentIdOfBlock - ) - }) - .filter(notUndefined) - - return new BlockNode( - this.activeScriptHash, - newChildren, - this.deltaBlock, - currentScriptRunId, - this.fragmentId, - this.deltaMsgReceivedAt - ) - } - - public getElements(elementSet?: Set): Set { - if (isNullOrUndefined(elementSet)) { - elementSet = new Set() - } - - for (const child of this.children) { - child.getElements(elementSet) - } - - return elementSet - } -} - -/** - * The root of our data tree. It contains the app's top-level BlockNodes. - */ -export class AppRoot { - readonly root: BlockNode - - /* The hash of the main script that creates this AppRoot. */ - readonly mainScriptHash: string - - /* The app logo, if it exists. */ - private appLogo: AppLogo | null - - /** - * Create an empty AppRoot with a placeholder "skeleton" element. - */ - public static empty( - mainScriptHash = "", - isInitialRender = true, - sidebarElements?: BlockNode, - logo?: Logo | null - ): AppRoot { - const mainNodes: AppNode[] = [] - - let waitElement: Element | undefined - - switch (getLoadingScreenType()) { - case LoadingScreenType.NONE: - break - - case LoadingScreenType.V1: - // Only show the v1 loading state when it's the initial render. - // This is how v1 used to work, and we don't want any backward - // incompatibility. - if (isInitialRender) { - waitElement = makeElementWithInfoText("Please wait...") - } - break - - default: - waitElement = makeAppSkeletonElement() - } - - if (waitElement) { - mainNodes.push( - new ElementNode( - waitElement, - ForwardMsgMetadata.create({}), - NO_SCRIPT_RUN_ID, - mainScriptHash - ) - ) - } - - const main = new BlockNode( - mainScriptHash, - mainNodes, - new BlockProto({ allowEmpty: true }), - NO_SCRIPT_RUN_ID - ) - - const sidebar = - sidebarElements || - new BlockNode( - mainScriptHash, - [], - new BlockProto({ allowEmpty: true }), - NO_SCRIPT_RUN_ID - ) - - const event = new BlockNode( - mainScriptHash, - [], - new BlockProto({ allowEmpty: true }), - NO_SCRIPT_RUN_ID - ) - - const bottom = new BlockNode( - mainScriptHash, - [], - new BlockProto({ allowEmpty: true }), - NO_SCRIPT_RUN_ID - ) - - // Persist logo between pages to avoid flicker (MPA V1 - Issue #8815) - const appLogo = logo - ? { - logo, - activeScriptHash: mainScriptHash, - scriptRunId: NO_SCRIPT_RUN_ID, - } - : null - - return new AppRoot( - mainScriptHash, - new BlockNode(mainScriptHash, [main, sidebar, event, bottom]), - appLogo - ) - } - - public constructor( - mainScriptHash: string, - root: BlockNode, - appLogo: AppLogo | null = null - ) { - this.mainScriptHash = mainScriptHash - this.root = root - this.appLogo = appLogo - - // Verify that our root node has exactly 4 children: a 'main' block, - // a 'sidebar' block, a `bottom` block and an 'event' block. - if ( - this.root.children.length !== 4 || - isNullOrUndefined(this.main) || - isNullOrUndefined(this.sidebar) || - isNullOrUndefined(this.event) || - isNullOrUndefined(this.bottom) - ) { - // eslint-disable-next-line @typescript-eslint/no-base-to-string, @typescript-eslint/restrict-template-expressions -- TODO: Fix this - throw new Error(`Invalid root node children! ${root}`) - } - } - - public get main(): BlockNode { - const [main] = this.root.children - return main as BlockNode - } - - public get sidebar(): BlockNode { - const [, sidebar] = this.root.children - return sidebar as BlockNode - } - - public get event(): BlockNode { - const [, , event] = this.root.children - return event as BlockNode - } - - public get bottom(): BlockNode { - const [, , , bottom] = this.root.children - return bottom as BlockNode - } - - public get logo(): Logo | null { - return this.appLogo?.logo ?? null - } - - public appRootWithLogo(logo: Logo, metadata: LogoMetadata): AppRoot { - return new AppRoot(this.mainScriptHash, this.root, { - logo, - ...metadata, - }) - } - - public applyDelta( - scriptRunId: string, - delta: Delta, - metadata: ForwardMsgMetadata - ): AppRoot { - // The full path to the AppNode within the element tree. - // Used to find and update the element node specified by this Delta. - const { deltaPath, activeScriptHash } = metadata - switch (delta.type) { - case "newElement": { - const element = delta.newElement as Element - return this.addElement( - deltaPath, - scriptRunId, - element, - metadata, - activeScriptHash, - delta.fragmentId - ) - } - - case "addBlock": { - const deltaMsgReceivedAt = Date.now() - return this.addBlock( - deltaPath, - delta.addBlock as BlockProto, - scriptRunId, - activeScriptHash, - delta.fragmentId, - deltaMsgReceivedAt - ) - } - - case "arrowAddRows": { - try { - return this.arrowAddRows( - deltaPath, - delta.arrowAddRows as ArrowNamedDataSet, - scriptRunId - ) - } catch (error) { - const errorElement = makeElementWithErrorText( - ensureError(error).message - ) - return this.addElement( - deltaPath, - scriptRunId, - errorElement, - metadata, - activeScriptHash - ) - } - } - - default: { - throw new Error(`Unrecognized deltaType: '${delta.type}'`) - } - } - } - - filterMainScriptElements(mainScriptHash: string): AppRoot { - // clears all nodes that are not associated with the mainScriptHash - // Get the current script run id from one of the children - const currentScriptRunId = this.main.scriptRunId - const main = - this.main.filterMainScriptElements(mainScriptHash) || - new BlockNode(mainScriptHash) - const sidebar = - this.sidebar.filterMainScriptElements(mainScriptHash) || - new BlockNode(mainScriptHash) - const event = - this.event.filterMainScriptElements(mainScriptHash) || - new BlockNode(mainScriptHash) - const bottom = - this.bottom.filterMainScriptElements(mainScriptHash) || - new BlockNode(mainScriptHash) - const appLogo = - this.appLogo?.activeScriptHash === mainScriptHash ? this.appLogo : null - - return new AppRoot( - mainScriptHash, - new BlockNode( - mainScriptHash, - [main, sidebar, event, bottom], - new BlockProto({ allowEmpty: true }), - currentScriptRunId - ), - appLogo - ) - } - - public clearStaleNodes( - currentScriptRunId: string, - fragmentIdsThisRun?: Array - ): AppRoot { - const main = - this.main.clearStaleNodes(currentScriptRunId, fragmentIdsThisRun) || - new BlockNode(this.mainScriptHash) - const sidebar = - this.sidebar.clearStaleNodes(currentScriptRunId, fragmentIdsThisRun) || - new BlockNode(this.mainScriptHash) - const event = - this.event.clearStaleNodes(currentScriptRunId, fragmentIdsThisRun) || - new BlockNode(this.mainScriptHash) - const bottom = - this.bottom.clearStaleNodes(currentScriptRunId, fragmentIdsThisRun) || - new BlockNode(this.mainScriptHash) - - // Check if we're running a fragment, ensure logo isn't cleared as stale (Issue #10350/#10382) - const isFragmentRun = fragmentIdsThisRun && fragmentIdsThisRun.length > 0 - const appLogo = - isFragmentRun || this.appLogo?.scriptRunId === currentScriptRunId - ? this.appLogo - : null - - return new AppRoot( - this.mainScriptHash, - new BlockNode( - this.mainScriptHash, - [main, sidebar, event, bottom], - new BlockProto({ allowEmpty: true }), - currentScriptRunId - ), - appLogo - ) - } - - /** Return a Set containing all Elements in the tree. */ - public getElements(): Set { - const elements = new Set() - this.main.getElements(elements) - this.sidebar.getElements(elements) - this.event.getElements(elements) - this.bottom.getElements(elements) - return elements - } - - private addElement( - deltaPath: number[], - scriptRunId: string, - element: Element, - metadata: ForwardMsgMetadata, - activeScriptHash: string, - fragmentId?: string - ): AppRoot { - const elementNode = new ElementNode( - element, - metadata, - scriptRunId, - activeScriptHash, - fragmentId - ) - return new AppRoot( - this.mainScriptHash, - this.root.setIn(deltaPath, elementNode, scriptRunId), - this.appLogo - ) - } - - private addBlock( - deltaPath: number[], - block: BlockProto, - scriptRunId: string, - activeScriptHash: string, - fragmentId?: string, - deltaMsgReceivedAt?: number - ): AppRoot { - const existingNode = this.root.getIn(deltaPath) - - // If we're replacing an existing Block of the same type, this new Block - // inherits the existing Block's children. This preserves two things: - // 1. Widget State - // 2. React state of all elements - let children: AppNode[] = [] - if ( - existingNode instanceof BlockNode && - existingNode.deltaBlock.type === block.type - ) { - children = existingNode.children - } - - const blockNode = new BlockNode( - activeScriptHash, - children, - block, - scriptRunId, - fragmentId, - deltaMsgReceivedAt - ) - return new AppRoot( - this.mainScriptHash, - this.root.setIn(deltaPath, blockNode, scriptRunId), - this.appLogo - ) - } - - private arrowAddRows( - deltaPath: number[], - namedDataSet: ArrowNamedDataSet, - scriptRunId: string - ): AppRoot { - const existingNode = this.root.getIn(deltaPath) as ElementNode - if (isNullOrUndefined(existingNode)) { - // eslint-disable-next-line @typescript-eslint/restrict-template-expressions - throw new Error(`Can't arrowAddRows: invalid deltaPath: ${deltaPath}`) - } - - const elementNode = existingNode.arrowAddRows(namedDataSet, scriptRunId) - return new AppRoot( - this.mainScriptHash, - this.root.setIn(deltaPath, elementNode, scriptRunId), - this.appLogo - ) - } -} - -/** Iterates over datasets and converts data to Quiver. */ -function wrapDatasets(datasets: IArrowNamedDataSet[]): WrappedNamedDataset[] { - return datasets.map((dataset: IArrowNamedDataSet) => { - return { - hasName: dataset.hasName as boolean, - name: dataset.name as string, - data: new Quiver(dataset.data as IArrow), - } - }) -} +export type { AppNode } from "./render-tree/AppNode.interface" +export { AppRoot } from "./render-tree/AppRoot" +export { BlockNode } from "./render-tree/BlockNode" +export { ElementNode } from "./render-tree/ElementNode" diff --git a/frontend/lib/src/render-tree/AppNode.interface.ts b/frontend/lib/src/render-tree/AppNode.interface.ts new file mode 100644 index 00000000000..41cc32a4f78 --- /dev/null +++ b/frontend/lib/src/render-tree/AppNode.interface.ts @@ -0,0 +1,122 @@ +/** + * Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2025) + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { Element } from "@streamlit/protobuf" + +/** + * The Generic ID of the script run this node was generated in. + */ +export const NO_SCRIPT_RUN_ID = "NO_SCRIPT_RUN_ID" + +/** + * An immutable node of the "App Data Tree". + * + * Trees are composed of `ElementNode` leaves, which contain data about + * a single visual element, and `BlockNode` branches, which determine the + * layout of a group of children nodes. + * + * A simple tree might look like this: + * + * AppRoot + * ├── BlockNode ("main") + * │ ├── ElementNode (text: "Ahoy, Streamlit!") + * │ └── ElementNode (button: "Don't Push This") + * └── BlockNode ("sidebar") + * └── ElementNode (checkbox: "Batten The Hatches") + * + * To build this tree, the frontend receives `Delta` messages from Python, + * each of which corresponds to a tree mutation ("add an element", + * "add a block", "add rows to an existing element"). The frontend builds the + * tree bit by bit in response to these `Delta`s. + * + * To render the app, the `AppView` class walks this tree and outputs + * a corresponding DOM structure, using React, that's essentially a mapping + * of `AppElement` -> `ReactNode`. This rendering happens "live" - that is, + * the app is re-rendered each time a new `Delta` is received. + * + * Because the app gets re-rendered frequently, updates need to be fast. + * Our React components - the building blocks of the app - are "pure" + * (see https://reactjs.org/docs/react-api.html#reactpurecomponent), which + * means that React uses shallow comparison to determine which ReactNodes to + * update. + * + * Thus, each node in our tree is _immutable_ - any change to a `AppNode` + * actually results in a *new* `AppNode` instance. This occurs recursively, + * so inserting a new `ElementNode` into the tree will also result in new + * `BlockNode`s for each of that Element's ancestors, all the way up to the + * root node. Then, when React re-renders the app, it will re-traverse the new + * nodes that have been created, and rebuild just the bits of the app that + * have changed. + */ +export interface AppNode { + /** + * The ID of the script run this node was generated in. When a script finishes + * running, the app prunes all stale nodes. + */ + readonly scriptRunId: string + + /** + * The ID of the fragment that sent the Delta creating this AppNode. If this + * AppNode was not created by a fragment, this field is falsy. + */ + readonly fragmentId?: string + + /** + * The hash of the script that created this node. + */ + readonly activeScriptHash?: string + + // A timestamp indicating based on which delta message the node was created. + // If the node was created without a delta message, this field is undefined. + // This helps us to update React components based on a new backend message even though other + // props have not changed; this can happen for UI-only interactions such as dismissing a dialog. + readonly deltaMsgReceivedAt?: number + + /** + * Return the AppNode for the given index path, or undefined if the path + * is invalid. + */ + getIn(path: number[]): AppNode | undefined + + /** + * Return a copy of this node with a new element set at the given index + * path. Throws an error if the path is invalid. + */ + setIn(path: number[], node: AppNode, scriptRunId: string): AppNode + + /** + * Recursively remove children nodes whose activeScriptHash is no longer + * associated with the mainScriptHash. + */ + filterMainScriptElements(mainScriptHash: string): AppNode | undefined + + /** + * Recursively remove children nodes whose scriptRunId is no longer current. + * If this node should no longer exist, return undefined. + */ + clearStaleNodes( + currentScriptRunId: string, + fragmentIdsThisRun?: Array, + fragmentIdOfBlock?: string + ): AppNode | undefined + + /** + * Return a Set of all the Elements contained in the tree. + * If an existing Set is passed in, that Set will be mutated and returned. + * Otherwise, a new Set will be created and will be returned. + */ + getElements(elementSet?: Set): Set +} diff --git a/frontend/lib/src/render-tree/AppRoot.test.ts b/frontend/lib/src/render-tree/AppRoot.test.ts new file mode 100644 index 00000000000..f807c4a8e62 --- /dev/null +++ b/frontend/lib/src/render-tree/AppRoot.test.ts @@ -0,0 +1,664 @@ +/** + * Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2025) + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { MockInstance } from "vitest" + +import { Delta as DeltaProto, Logo as LogoProto } from "@streamlit/protobuf" + +import { NO_SCRIPT_RUN_ID } from "./AppNode.interface" +import { AppRoot } from "./AppRoot" +import { BlockNode } from "./BlockNode" +import { ElementNode } from "./ElementNode" +import { + block, + FAKE_SCRIPT_HASH, + forwardMsgMetadata, + makeProto, + text, +} from "./test-utils" + +// prettier-ignore +const BLOCK = block([ + text("1"), + block([ + text("2"), + ]), +]) + +// Initialize new AppRoot with a main block node and three child block nodes - sidebar, events and bottom. +const ROOT = new AppRoot( + FAKE_SCRIPT_HASH, + new BlockNode(FAKE_SCRIPT_HASH, [ + BLOCK, + new BlockNode(FAKE_SCRIPT_HASH), + new BlockNode(FAKE_SCRIPT_HASH), + new BlockNode(FAKE_SCRIPT_HASH), + ]) +) + +describe("AppRoot", () => { + describe("AppRoot.empty", () => { + let windowSpy: MockInstance + + beforeEach(() => { + windowSpy = vi.spyOn(window, "window", "get") + }) + + afterEach(() => { + windowSpy.mockRestore() + }) + + it("creates empty tree except for a skeleton", () => { + windowSpy.mockImplementation(() => ({ + location: { + search: "", + }, + })) + const empty = AppRoot.empty(FAKE_SCRIPT_HASH) + + expect(empty.main.children.length).toBe(1) + const child = empty.main.getIn([0]) as ElementNode + expect(child.element.skeleton).not.toBeNull() + + expect(empty.sidebar.isEmpty).toBe(true) + }) + + it("sets the main script hash and active script hash", () => { + windowSpy.mockImplementation(() => ({ + location: { + search: "", + }, + })) + const empty = AppRoot.empty(FAKE_SCRIPT_HASH) + + expect(empty.mainScriptHash).toBe(FAKE_SCRIPT_HASH) + expect(empty.main.activeScriptHash).toBe(FAKE_SCRIPT_HASH) + expect(empty.sidebar.activeScriptHash).toBe(FAKE_SCRIPT_HASH) + expect(empty.event.activeScriptHash).toBe(FAKE_SCRIPT_HASH) + expect(empty.bottom.activeScriptHash).toBe(FAKE_SCRIPT_HASH) + expect(empty.root.activeScriptHash).toBe(FAKE_SCRIPT_HASH) + }) + + it("creates empty tree with no loading screen if query param is set", () => { + windowSpy.mockImplementation(() => ({ + location: { + search: "?embed_options=hide_loading_screen", + }, + })) + + const empty = AppRoot.empty(FAKE_SCRIPT_HASH) + + expect(empty.main.isEmpty).toBe(true) + expect(empty.sidebar.isEmpty).toBe(true) + }) + + it("creates empty tree with v1 loading screen if query param is set", () => { + windowSpy.mockImplementation(() => ({ + location: { + search: "?embed_options=show_loading_screen_v1", + }, + })) + + const empty = AppRoot.empty(FAKE_SCRIPT_HASH) + + expect(empty.main.children.length).toBe(1) + const child = empty.main.getIn([0]) as ElementNode + expect(child.element.alert).toBeDefined() + + expect(empty.sidebar.isEmpty).toBe(true) + }) + + it("creates empty tree with v2 loading screen if query param is set", () => { + windowSpy.mockImplementation(() => ({ + location: { + search: "?embed_options=show_loading_screen_v2", + }, + })) + + const empty = AppRoot.empty(FAKE_SCRIPT_HASH) + + expect(empty.main.children.length).toBe(1) + const child = empty.main.getIn([0]) as ElementNode + expect(child.element.skeleton).not.toBeNull() + + expect(empty.sidebar.isEmpty).toBe(true) + }) + + it("creates empty tree with no loading screen if query param is v1 and it's not first load", () => { + windowSpy.mockImplementation(() => ({ + location: { + search: "?embed_options=show_loading_screen_v1", + }, + })) + + const empty = AppRoot.empty(FAKE_SCRIPT_HASH, false) + + expect(empty.main.isEmpty).toBe(true) + expect(empty.sidebar.isEmpty).toBe(true) + }) + + it("passes logo to new Root if empty is called with logo", () => { + windowSpy.mockImplementation(() => ({ + location: { + search: "", + }, + })) + const logo = LogoProto.create({ + image: + "https://global.discourse-cdn.com/business7/uploads/streamlit/original/2X/8/8cb5b6c0e1fe4e4ebfd30b769204c0d30c332fec.png", + }) + + // Replicate .empty call on initial render + const empty = AppRoot.empty("", true) + expect(empty.logo).toBeNull() + + // Replicate .empty call in AppNav's clearPageElements for MPA V1 + const empty2 = AppRoot.empty(FAKE_SCRIPT_HASH, false, undefined, logo) + expect(empty2.logo).not.toBeNull() + }) + }) + + describe("AppRoot.filterMainScriptElements", () => { + it("does not clear nodes associated with main script hash", () => { + // Add a new element and clear stale nodes + const delta = makeProto(DeltaProto, { + newElement: { text: { body: "newElement!" } }, + }) + const newRoot = ROOT.applyDelta( + "new_session_id", + delta, + forwardMsgMetadata([0, 1, 1]) + ).filterMainScriptElements(FAKE_SCRIPT_HASH) + + // We should now only have a single element, inside a single block + expect(newRoot.main.getIn([1, 1])).toBeTextNode("newElement!") + expect(newRoot.getElements().size).toBe(3) + }) + + it("clears nodes not associated with main script hash", () => { + // Add a new element and clear stale nodes + const delta = makeProto(DeltaProto, { + newElement: { text: { body: "newElement!" } }, + }) + const newRoot = ROOT.applyDelta( + "new_session_id", + delta, + forwardMsgMetadata([0, 1, 1], "DIFFERENT_HASH") + ).filterMainScriptElements(FAKE_SCRIPT_HASH) + + // We should now only have a single element, inside a single block + expect(newRoot.main.getIn([1, 1])).toBeUndefined() + expect(newRoot.getElements().size).toBe(2) + }) + }) + + describe("AppRoot.applyDelta", () => { + it("handles 'newElement' deltas", () => { + const delta = makeProto(DeltaProto, { + newElement: { text: { body: "newElement!" } }, + }) + const newRoot = ROOT.applyDelta( + "new_session_id", + delta, + forwardMsgMetadata([0, 1, 1]) + ) + + const newNode = newRoot.main.getIn([1, 1]) as ElementNode + expect(newNode).toBeTextNode("newElement!") + + // Check that our new scriptRunId has been set only on the touched nodes + expect(newRoot.main.scriptRunId).toBe("new_session_id") + expect(newRoot.main.fragmentId).toBe(undefined) + expect(newRoot.main.deltaMsgReceivedAt).toBe(undefined) + expect(newRoot.main.getIn([0])?.scriptRunId).toBe(NO_SCRIPT_RUN_ID) + expect(newRoot.main.getIn([1])?.scriptRunId).toBe("new_session_id") + expect(newRoot.main.getIn([1, 0])?.scriptRunId).toBe(NO_SCRIPT_RUN_ID) + expect(newRoot.main.getIn([1, 1])?.scriptRunId).toBe("new_session_id") + expect(newNode.activeScriptHash).toBe(FAKE_SCRIPT_HASH) + expect(newRoot.sidebar.scriptRunId).toBe(NO_SCRIPT_RUN_ID) + }) + + it("handles 'addBlock' deltas", () => { + const delta = makeProto(DeltaProto, { addBlock: {} }) + const newRoot = ROOT.applyDelta( + "new_session_id", + delta, + forwardMsgMetadata([0, 1, 1]) + ) + + const newNode = newRoot.main.getIn([1, 1]) as BlockNode + expect(newNode).toBeDefined() + + // Check that our new scriptRunId has been set only on the touched nodes + expect(newRoot.main.scriptRunId).toBe("new_session_id") + expect(newRoot.main.fragmentId).toBe(undefined) + expect(newRoot.main.deltaMsgReceivedAt).toBe(undefined) + expect(newRoot.main.getIn([0])?.scriptRunId).toBe(NO_SCRIPT_RUN_ID) + expect(newRoot.main.getIn([1])?.scriptRunId).toBe("new_session_id") + expect(newRoot.main.getIn([1, 0])?.scriptRunId).toBe(NO_SCRIPT_RUN_ID) + expect(newRoot.main.getIn([1, 1])?.scriptRunId).toBe("new_session_id") + expect(newNode.activeScriptHash).toBe(FAKE_SCRIPT_HASH) + expect(newRoot.sidebar.scriptRunId).toBe(NO_SCRIPT_RUN_ID) + }) + + it("removes a block's children if the block type changes for the same delta path", () => { + const newRoot = ROOT.applyDelta( + "script_run_id", + makeProto(DeltaProto, { + addBlock: { + expandable: { + expanded: true, + label: "label", + icon: "", + }, + }, + }), + forwardMsgMetadata([0, 1, 1]) + ).applyDelta( + "script_run_id", + makeProto(DeltaProto, { + newElement: { text: { body: "newElement!" } }, + }), + forwardMsgMetadata([0, 1, 1, 0]) + ) + + const newNode = newRoot.main.getIn([1, 1]) as BlockNode + expect(newNode).toBeDefined() + expect(newNode.deltaBlock.type).toBe("expandable") + expect(newNode.children.length).toBe(1) + + const newRoot2 = newRoot.applyDelta( + "new_script_run_id", + makeProto(DeltaProto, { + addBlock: { + tabContainer: {}, + }, + }), + forwardMsgMetadata([0, 1, 1]) + ) + + const replacedBlock = newRoot2.main.getIn([1, 1]) as BlockNode + expect(replacedBlock).toBeDefined() + expect(replacedBlock.deltaBlock.type).toBe("tabContainer") + expect(replacedBlock.children.length).toBe(0) + }) + + it("will not remove a block's children if the block type is the same for the same delta path", () => { + const newRoot = ROOT.applyDelta( + "script_run_id", + makeProto(DeltaProto, { + addBlock: { + expandable: { + expanded: true, + label: "label", + icon: "", + }, + }, + }), + forwardMsgMetadata([0, 1, 1]) + ).applyDelta( + "script_run_id", + makeProto(DeltaProto, { + newElement: { text: { body: "newElement!" } }, + }), + forwardMsgMetadata([0, 1, 1, 0]) + ) + + const newNode = newRoot.main.getIn([1, 1]) as BlockNode + expect(newNode).toBeDefined() + expect(newNode.deltaBlock.type).toBe("expandable") + expect(newNode.children.length).toBe(1) + + const newRoot2 = newRoot.applyDelta( + "new_script_run_id", + makeProto(DeltaProto, { + addBlock: { + expandable: { + expanded: true, + label: "other label", + icon: "", + }, + }, + }), + forwardMsgMetadata([0, 1, 1]) + ) + + const replacedBlock = newRoot2.main.getIn([1, 1]) as BlockNode + expect(replacedBlock).toBeDefined() + expect(replacedBlock.deltaBlock.type).toBe("expandable") + expect(replacedBlock.children.length).toBe(1) + }) + + it("specifies active script hash on 'newElement' deltas", () => { + const delta = makeProto(DeltaProto, { + newElement: { text: { body: "newElement!" } }, + }) + const NEW_FAKE_SCRIPT_HASH = "new_fake_script_hash" + const newRoot = ROOT.applyDelta( + "new_session_id", + delta, + forwardMsgMetadata([0, 1, 1], NEW_FAKE_SCRIPT_HASH) + ) + + const newNode = newRoot.main.getIn([1, 1]) as ElementNode + expect(newNode).toBeDefined() + + // Check that our new other nodes are not affected by the new script hash + expect(newRoot.main.getIn([1, 0])?.activeScriptHash).toBe( + FAKE_SCRIPT_HASH + ) + expect(newNode.activeScriptHash).toBe(NEW_FAKE_SCRIPT_HASH) + }) + + it("specifies active script hash on 'addBlock' deltas", () => { + const delta = makeProto(DeltaProto, { addBlock: {} }) + const NEW_FAKE_SCRIPT_HASH = "new_fake_script_hash" + const newRoot = ROOT.applyDelta( + "new_session_id", + delta, + forwardMsgMetadata([0, 1, 1], NEW_FAKE_SCRIPT_HASH) + ) + + const newNode = newRoot.main.getIn([1, 1]) as BlockNode + expect(newNode).toBeDefined() + + // Check that our new scriptRunId has been set only on the touched nodes + expect(newRoot.main.getIn([1, 0])?.activeScriptHash).toBe( + FAKE_SCRIPT_HASH + ) + expect(newNode.activeScriptHash).toBe(NEW_FAKE_SCRIPT_HASH) + }) + + it("can set fragmentId in 'newElement' deltas", () => { + const delta = makeProto(DeltaProto, { + newElement: { text: { body: "newElement!" } }, + fragmentId: "myFragmentId", + }) + const newRoot = ROOT.applyDelta( + "new_session_id", + delta, + forwardMsgMetadata([0, 1, 1]) + ) + + const newNode = newRoot.main.getIn([1, 1]) as ElementNode + expect(newNode.fragmentId).toBe("myFragmentId") + }) + + it("can set fragmentId in 'addBlock' deltas", () => { + const delta = makeProto(DeltaProto, { + addBlock: {}, + fragmentId: "myFragmentId", + }) + const newRoot = ROOT.applyDelta( + "new_session_id", + delta, + forwardMsgMetadata([0, 1, 1]) + ) + + const newNode = newRoot.main.getIn([1, 1]) as BlockNode + expect(newNode.fragmentId).toBe("myFragmentId") + }) + + it("timestamp is set on BlockNode as message id", () => { + const timestamp = new Date(Date.UTC(2017, 1, 14)).valueOf() + Date.now = vi.fn(() => timestamp) + const delta = makeProto(DeltaProto, { + addBlock: {}, + }) + const newRoot = ROOT.applyDelta( + "new_session_id", + delta, + forwardMsgMetadata([0, 1, 1]) + ) + + const newNode = newRoot.main.getIn([1, 1]) as BlockNode + expect(newNode.deltaMsgReceivedAt).toBe(timestamp) + }) + }) + + describe("AppRoot.clearStaleNodes", () => { + it("clears stale nodes", () => { + // Add a new element and clear stale nodes + const delta = makeProto(DeltaProto, { + newElement: { text: { body: "newElement!" } }, + }) + const newRoot = ROOT.applyDelta( + "new_session_id", + delta, + forwardMsgMetadata([0, 1, 1]) + ).clearStaleNodes("new_session_id", []) + + // We should now only have a single element, inside a single block + expect(newRoot.main.getIn([0, 0])).toBeTextNode("newElement!") + expect(newRoot.getElements().size).toBe(1) + }) + + it("clears a stale logo", () => { + const logo = LogoProto.create({ + image: + "https://global.discourse-cdn.com/business7/uploads/streamlit/original/2X/8/8cb5b6c0e1fe4e4ebfd30b769204c0d30c332fec.png", + }) + const newRoot = ROOT.appRootWithLogo(logo, { + activeScriptHash: "hash", + scriptRunId: "script_run_id", + }) + expect(newRoot.logo).not.toBeNull() + + const newNewRoot = newRoot.clearStaleNodes("new_script_run_id", []) + expect(newNewRoot.logo).toBeNull() + }) + + it("does not clear logo on fragment run", () => { + const logo = LogoProto.create({ + image: + "https://global.discourse-cdn.com/business7/uploads/streamlit/original/2X/8/8cb5b6c0e1fe4e4ebfd30b769204c0d30c332fec.png", + }) + const newRoot = ROOT.appRootWithLogo(logo, { + activeScriptHash: "hash", + scriptRunId: "script_run_id", + }) + expect(newRoot.logo).not.toBeNull() + + const newNewRoot = newRoot.clearStaleNodes("new_script_run_id", [ + "my_fragment_id", + ]) + expect(newNewRoot.logo).not.toBeNull() + }) + + it("handles currentFragmentId correctly", () => { + const tabContainerProto = makeProto(DeltaProto, { + addBlock: { tabContainer: {}, allowEmpty: false }, + fragmentId: "my_fragment_id", + }) + const tab1 = makeProto(DeltaProto, { + addBlock: { tab: { label: "tab1" }, allowEmpty: true }, + fragmentId: "my_fragment_id", + }) + const tab2 = makeProto(DeltaProto, { + addBlock: { tab: { label: "tab2" }, allowEmpty: true }, + fragmentId: "my_fragment_id", + }) + + // const BLOCK = block([text("1"), block([text("2")])]) + const root = AppRoot.empty(FAKE_SCRIPT_HASH) + // Block not corresponding to my_fragment_id. Should be preserved. + .applyDelta( + "old_session_id", + makeProto(DeltaProto, { addBlock: { allowEmpty: true } }), + forwardMsgMetadata([0, 0]) + ) + // Element in block unrelated to my_fragment_id. Should be preserved. + .applyDelta( + "old_session_id", + makeProto(DeltaProto, { + newElement: { text: { body: "oldElement!" } }, + }), + forwardMsgMetadata([0, 0, 0]) + ) + // Another element in block unrelated to my_fragment_id. Should be preserved. + .applyDelta( + "old_session_id", + makeProto(DeltaProto, { + newElement: { text: { body: "oldElement2!" } }, + fragmentId: "other_fragment_id", + }), + forwardMsgMetadata([0, 0, 1]) + ) + // Old element related to my_fragment_id but in an unrelated block. Should be preserved. + .applyDelta( + "old_session_id", + makeProto(DeltaProto, { + newElement: { text: { body: "oldElement4!" } }, + fragmentId: "my_fragment_id", + }), + forwardMsgMetadata([0, 0, 2]) + ) + // Block corresponding to my_fragment_id + .applyDelta( + "new_session_id", + makeProto(DeltaProto, { + addBlock: { allowEmpty: false }, + fragmentId: "my_fragment_id", + }), + forwardMsgMetadata([0, 1]) + ) + // Old element related to my_fragment_id. Should be pruned. + .applyDelta( + "old_session_id", + makeProto(DeltaProto, { + newElement: { text: { body: "oldElement3!" } }, + fragmentId: "my_fragment_id", + }), + forwardMsgMetadata([0, 1, 0]) + ) + // New element related to my_fragment_id. Should be preserved. + .applyDelta( + "new_session_id", + makeProto(DeltaProto, { + newElement: { text: { body: "newElement!" } }, + fragmentId: "my_fragment_id", + }), + forwardMsgMetadata([0, 1, 1]) + ) + // New element container related to my_fragment_id, having children which will be handled individually + // Create a tab container with two tabs in the old session; then send new delta with the container and + // only one tab. The second tab with the old_session_id should be pruned. + .applyDelta( + "old_session_id", + tabContainerProto, + forwardMsgMetadata([0, 2]) + ) + .applyDelta("old_session_id", tab1, forwardMsgMetadata([0, 2, 0])) + .applyDelta("old_session_id", tab2, forwardMsgMetadata([0, 2, 1])) + .applyDelta( + "new_session_id", + tabContainerProto, + forwardMsgMetadata([0, 2]) + ) + .applyDelta("new_session_id", tab1, forwardMsgMetadata([0, 2, 0])) + + const pruned = root.clearStaleNodes("new_session_id", ["my_fragment_id"]) + + expect(pruned.main.getIn([0])).toBeInstanceOf(BlockNode) + expect((pruned.main.getIn([0]) as BlockNode).children).toHaveLength(3) + expect(pruned.main.getIn([0, 0])).toBeTextNode("oldElement!") + expect(pruned.main.getIn([0, 1])).toBeTextNode("oldElement2!") + expect(pruned.main.getIn([0, 2])).toBeTextNode("oldElement4!") + + expect(pruned.main.getIn([1])).toBeInstanceOf(BlockNode) + expect((pruned.main.getIn([1]) as BlockNode).children).toHaveLength(1) + expect(pruned.main.getIn([1, 0])).toBeTextNode("newElement!") + + expect(pruned.main.getIn([2])).toBeInstanceOf(BlockNode) + expect((pruned.main.getIn([2]) as BlockNode).children).toHaveLength(1) + expect( + (pruned.main.getIn([2, 0]) as BlockNode).deltaBlock.tab?.label + ).toContain("tab1") + }) + + it("clear childNodes of a block node in fragment run", () => { + // Add a new element and clear stale nodes + const delta = makeProto(DeltaProto, { + newElement: { text: { body: "newElement!" } }, + fragmentId: "my_fragment_id", + }) + const newRoot = AppRoot.empty(FAKE_SCRIPT_HASH) + // Block corresponding to my_fragment_id + .applyDelta( + "new_session_id", + makeProto(DeltaProto, { + addBlock: { vertical: {}, allowEmpty: false }, + fragmentId: "my_fragment_id", + }), + forwardMsgMetadata([0, 0]) + ) + .applyDelta("new_session_id", delta, forwardMsgMetadata([0, 0, 0])) + // Block with child where scriptRunId is different + .applyDelta( + "new_session_id", + makeProto(DeltaProto, { + addBlock: { vertical: {}, allowEmpty: false }, + fragmentId: "my_fragment_id", + }), + forwardMsgMetadata([0, 1]) + ) + .applyDelta("new_session_id", delta, forwardMsgMetadata([0, 1, 0])) + .applyDelta("new_session_id", delta, forwardMsgMetadata([0, 1, 1])) + // this child is a nested fragment_id from an old run and should be pruned + .applyDelta( + "old_session_id", + makeProto(DeltaProto, { + newElement: { text: { body: "oldElement!" } }, + fragmentId: "my_nested_fragment_id", + }), + forwardMsgMetadata([0, 1, 2]) + ) + // this child is a nested fragment_id from the same run and should be preserved + .applyDelta( + "new_session_id", + makeProto(DeltaProto, { + newElement: { text: { body: "newElement!" } }, + fragmentId: "my_nested_fragment_id", + }), + forwardMsgMetadata([0, 1, 3]) + ) + + expect((newRoot.main.getIn([1]) as BlockNode).children).toHaveLength(4) + + const pruned = newRoot.clearStaleNodes("new_session_id", [ + "my_fragment_id", + ]) + + expect(pruned.main.getIn([0])).toBeInstanceOf(BlockNode) + expect((pruned.main.getIn([0]) as BlockNode).children).toHaveLength(1) + expect(pruned.main.getIn([1])).toBeInstanceOf(BlockNode) + // the stale nested fragment child should have been pruned + expect((pruned.main.getIn([1]) as BlockNode).children).toHaveLength(3) + }) + }) + + describe("AppRoot.getElements", () => { + it("returns all elements", () => { + // We have elements at main.[0] and main.[1, 0] + expect(ROOT.getElements()).toEqual( + new Set([ + (ROOT.main.getIn([0]) as ElementNode).element, + (ROOT.main.getIn([1, 0]) as ElementNode).element, + ]) + ) + }) + }) +}) diff --git a/frontend/lib/src/render-tree/AppRoot.ts b/frontend/lib/src/render-tree/AppRoot.ts new file mode 100644 index 00000000000..b22bff96e12 --- /dev/null +++ b/frontend/lib/src/render-tree/AppRoot.ts @@ -0,0 +1,426 @@ +/** + * Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2025) + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { + ArrowNamedDataSet, + Block as BlockProto, + Delta, + Element, + ForwardMsgMetadata, + Logo, +} from "@streamlit/protobuf" + +import { ensureError } from "~lib/util/ErrorHandling" +import { + getLoadingScreenType, + isNullOrUndefined, + LoadingScreenType, + makeAppSkeletonElement, + makeElementWithErrorText, + makeElementWithInfoText, +} from "~lib/util/utils" + +import { AppNode } from "./AppNode.interface" +import { BlockNode } from "./BlockNode" +import { ElementNode } from "./ElementNode" + +const NO_SCRIPT_RUN_ID = "NO_SCRIPT_RUN_ID" + +interface LogoMetadata { + // Associated scriptHash that created the logo + activeScriptHash: string + + // Associated scriptRunId that created the logo + scriptRunId: string +} +interface AppLogo extends LogoMetadata { + logo: Logo +} + +/** + * The root of our data tree. It contains the app's top-level BlockNodes. + */ +export class AppRoot { + readonly root: BlockNode + + /* The hash of the main script that creates this AppRoot. */ + readonly mainScriptHash: string + + /* The app logo, if it exists. */ + private appLogo: AppLogo | null + + /** + * Create an empty AppRoot with a placeholder "skeleton" element. + * @param mainScriptHash - The hash of the main script that creates this AppRoot. + * @param isInitialRender - Whether this is the initial render. + * @param sidebarElements - The elements to add to the sidebar (this was a relic + * of MPA V1 to maintain the sidebar from flickering, we don't use it anymore). + * @param logo - The logo to add to the app. + * @returns A new AppRoot with the given parameters. + */ + public static empty( + mainScriptHash = "", + isInitialRender = true, + sidebarElements?: BlockNode, + logo?: Logo | null + ): AppRoot { + const mainNodes: AppNode[] = [] + + let waitElement: Element | undefined + + switch (getLoadingScreenType()) { + case LoadingScreenType.NONE: + break + + case LoadingScreenType.V1: + // Only show the v1 loading state when it's the initial render. + // This is how v1 used to work, and we don't want any backward + // incompatibility. + if (isInitialRender) { + waitElement = makeElementWithInfoText("Please wait...") + } + break + + default: + waitElement = makeAppSkeletonElement() + } + + if (waitElement) { + mainNodes.push( + new ElementNode( + waitElement, + ForwardMsgMetadata.create({}), + NO_SCRIPT_RUN_ID, + mainScriptHash + ) + ) + } + + const main = new BlockNode( + mainScriptHash, + mainNodes, + new BlockProto({ allowEmpty: true }), + NO_SCRIPT_RUN_ID + ) + + const sidebar = + sidebarElements || + new BlockNode( + mainScriptHash, + [], + new BlockProto({ allowEmpty: true }), + NO_SCRIPT_RUN_ID + ) + + const event = new BlockNode( + mainScriptHash, + [], + new BlockProto({ allowEmpty: true }), + NO_SCRIPT_RUN_ID + ) + + const bottom = new BlockNode( + mainScriptHash, + [], + new BlockProto({ allowEmpty: true }), + NO_SCRIPT_RUN_ID + ) + + // Persist logo between pages to avoid flicker (MPA V1 - Issue #8815) + const appLogo = logo + ? { + logo, + activeScriptHash: mainScriptHash, + scriptRunId: NO_SCRIPT_RUN_ID, + } + : null + + return new AppRoot( + mainScriptHash, + new BlockNode(mainScriptHash, [main, sidebar, event, bottom]), + appLogo + ) + } + + public constructor( + mainScriptHash: string, + root: BlockNode, + appLogo: AppLogo | null = null + ) { + this.mainScriptHash = mainScriptHash + this.root = root + this.appLogo = appLogo + + // Verify that our root node has exactly 4 children: a 'main' block, + // a 'sidebar' block, a `bottom` block and an 'event' block. + if ( + this.root.children.length !== 4 || + isNullOrUndefined(this.main) || + isNullOrUndefined(this.sidebar) || + isNullOrUndefined(this.event) || + isNullOrUndefined(this.bottom) + ) { + // eslint-disable-next-line @typescript-eslint/no-base-to-string, @typescript-eslint/restrict-template-expressions -- TODO: Fix this + throw new Error(`Invalid root node children! ${root}`) + } + } + + public get main(): BlockNode { + const [main] = this.root.children + return main as BlockNode + } + + public get sidebar(): BlockNode { + const [, sidebar] = this.root.children + return sidebar as BlockNode + } + + public get event(): BlockNode { + const [, , event] = this.root.children + return event as BlockNode + } + + public get bottom(): BlockNode { + const [, , , bottom] = this.root.children + return bottom as BlockNode + } + + public get logo(): Logo | null { + return this.appLogo?.logo ?? null + } + + public appRootWithLogo(logo: Logo, metadata: LogoMetadata): AppRoot { + return new AppRoot(this.mainScriptHash, this.root, { + logo, + ...metadata, + }) + } + + public applyDelta( + scriptRunId: string, + delta: Delta, + metadata: ForwardMsgMetadata + ): AppRoot { + // The full path to the AppNode within the element tree. + // Used to find and update the element node specified by this Delta. + const { deltaPath, activeScriptHash } = metadata + switch (delta.type) { + case "newElement": { + const element = delta.newElement as Element + return this.addElement( + deltaPath, + scriptRunId, + element, + metadata, + activeScriptHash, + delta.fragmentId + ) + } + + case "addBlock": { + const deltaMsgReceivedAt = Date.now() + return this.addBlock( + deltaPath, + delta.addBlock as BlockProto, + scriptRunId, + activeScriptHash, + delta.fragmentId, + deltaMsgReceivedAt + ) + } + + case "arrowAddRows": { + try { + return this.arrowAddRows( + deltaPath, + delta.arrowAddRows as ArrowNamedDataSet, + scriptRunId + ) + } catch (error) { + const errorElement = makeElementWithErrorText( + ensureError(error).message + ) + return this.addElement( + deltaPath, + scriptRunId, + errorElement, + metadata, + activeScriptHash + ) + } + } + + default: { + throw new Error(`Unrecognized deltaType: '${delta.type}'`) + } + } + } + + filterMainScriptElements(mainScriptHash: string): AppRoot { + // clears all nodes that are not associated with the mainScriptHash + // Get the current script run id from one of the children + const currentScriptRunId = this.main.scriptRunId + const main = + this.main.filterMainScriptElements(mainScriptHash) || + new BlockNode(mainScriptHash) + const sidebar = + this.sidebar.filterMainScriptElements(mainScriptHash) || + new BlockNode(mainScriptHash) + const event = + this.event.filterMainScriptElements(mainScriptHash) || + new BlockNode(mainScriptHash) + const bottom = + this.bottom.filterMainScriptElements(mainScriptHash) || + new BlockNode(mainScriptHash) + const appLogo = + this.appLogo?.activeScriptHash === mainScriptHash ? this.appLogo : null + + return new AppRoot( + mainScriptHash, + new BlockNode( + mainScriptHash, + [main, sidebar, event, bottom], + new BlockProto({ allowEmpty: true }), + currentScriptRunId + ), + appLogo + ) + } + + public clearStaleNodes( + currentScriptRunId: string, + fragmentIdsThisRun?: Array + ): AppRoot { + const main = + this.main.clearStaleNodes(currentScriptRunId, fragmentIdsThisRun) || + new BlockNode(this.mainScriptHash) + const sidebar = + this.sidebar.clearStaleNodes(currentScriptRunId, fragmentIdsThisRun) || + new BlockNode(this.mainScriptHash) + const event = + this.event.clearStaleNodes(currentScriptRunId, fragmentIdsThisRun) || + new BlockNode(this.mainScriptHash) + const bottom = + this.bottom.clearStaleNodes(currentScriptRunId, fragmentIdsThisRun) || + new BlockNode(this.mainScriptHash) + + // Check if we're running a fragment, ensure logo isn't cleared as stale (Issue #10350/#10382) + const isFragmentRun = fragmentIdsThisRun && fragmentIdsThisRun.length > 0 + const appLogo = + isFragmentRun || this.appLogo?.scriptRunId === currentScriptRunId + ? this.appLogo + : null + + return new AppRoot( + this.mainScriptHash, + new BlockNode( + this.mainScriptHash, + [main, sidebar, event, bottom], + new BlockProto({ allowEmpty: true }), + currentScriptRunId + ), + appLogo + ) + } + + /** Return a Set containing all Elements in the tree. */ + public getElements(): Set { + const elements = new Set() + this.main.getElements(elements) + this.sidebar.getElements(elements) + this.event.getElements(elements) + this.bottom.getElements(elements) + return elements + } + + private addElement( + deltaPath: number[], + scriptRunId: string, + element: Element, + metadata: ForwardMsgMetadata, + activeScriptHash: string, + fragmentId?: string + ): AppRoot { + const elementNode = new ElementNode( + element, + metadata, + scriptRunId, + activeScriptHash, + fragmentId + ) + return new AppRoot( + this.mainScriptHash, + this.root.setIn(deltaPath, elementNode, scriptRunId), + this.appLogo + ) + } + + private addBlock( + deltaPath: number[], + block: BlockProto, + scriptRunId: string, + activeScriptHash: string, + fragmentId?: string, + deltaMsgReceivedAt?: number + ): AppRoot { + const existingNode = this.root.getIn(deltaPath) + + // If we're replacing an existing Block of the same type, this new Block + // inherits the existing Block's children. This preserves two things: + // 1. Widget State + // 2. React state of all elements + let children: AppNode[] = [] + if ( + existingNode instanceof BlockNode && + existingNode.deltaBlock.type === block.type + ) { + children = existingNode.children + } + + const blockNode = new BlockNode( + activeScriptHash, + children, + block, + scriptRunId, + fragmentId, + deltaMsgReceivedAt + ) + return new AppRoot( + this.mainScriptHash, + this.root.setIn(deltaPath, blockNode, scriptRunId), + this.appLogo + ) + } + + private arrowAddRows( + deltaPath: number[], + namedDataSet: ArrowNamedDataSet, + scriptRunId: string + ): AppRoot { + const existingNode = this.root.getIn(deltaPath) as ElementNode + if (isNullOrUndefined(existingNode)) { + // eslint-disable-next-line @typescript-eslint/restrict-template-expressions + throw new Error(`Can't arrowAddRows: invalid deltaPath: ${deltaPath}`) + } + + const elementNode = existingNode.arrowAddRows(namedDataSet, scriptRunId) + return new AppRoot( + this.mainScriptHash, + this.root.setIn(deltaPath, elementNode, scriptRunId), + this.appLogo + ) + } +} diff --git a/frontend/lib/src/render-tree/BlockNode.test.ts b/frontend/lib/src/render-tree/BlockNode.test.ts new file mode 100644 index 00000000000..9cd13953b42 --- /dev/null +++ b/frontend/lib/src/render-tree/BlockNode.test.ts @@ -0,0 +1,74 @@ +/** + * Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2025) + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { NO_SCRIPT_RUN_ID } from "./AppNode.interface" +import { block, text } from "./test-utils" + +// prettier-ignore +const BLOCK = block([ + text("1"), + block([ + text("2"), + ]), +]) + +describe("BlockNode", () => { + describe("BlockNode.getIn", () => { + it("handles shallow paths", () => { + const node = BLOCK.getIn([0]) + expect(node).toBeTextNode("1") + }) + + it("handles deep paths", () => { + const node = BLOCK.getIn([1, 0]) + expect(node).toBeTextNode("2") + }) + + it("returns undefined for invalid paths", () => { + const node = BLOCK.getIn([2, 3, 4]) + expect(node).toBeUndefined() + }) + }) + + describe("BlockNode.setIn", () => { + it("handles shallow paths", () => { + const newBlock = BLOCK.setIn([0], text("new"), NO_SCRIPT_RUN_ID) + expect(newBlock.getIn([0])).toBeTextNode("new") + + // Check BLOCK..newBlock diff is as expected. + expect(newBlock).not.toStrictEqual(BLOCK) + expect(newBlock.getIn([1])).toStrictEqual(BLOCK.getIn([1])) + }) + + it("handles deep paths", () => { + const newBlock = BLOCK.setIn([1, 1], text("new"), NO_SCRIPT_RUN_ID) + expect(newBlock.getIn([1, 1])).toBeTextNode("new") + + // Check BLOCK..newBlock diff is as expected + expect(newBlock).not.toStrictEqual(BLOCK) + expect(newBlock.getIn([0])).toStrictEqual(BLOCK.getIn([0])) + expect(newBlock.getIn([1])).not.toStrictEqual(BLOCK.getIn([1])) + expect(newBlock.getIn([1, 0])).toStrictEqual(BLOCK.getIn([1, 0])) + expect(newBlock.getIn([1, 1])).not.toStrictEqual(BLOCK.getIn([1, 1])) + }) + + it("throws an error for invalid paths", () => { + expect(() => BLOCK.setIn([1, 2], text("new"), NO_SCRIPT_RUN_ID)).toThrow( + "Bad 'setIn' index 2 (should be between [0, 1])" + ) + }) + }) +}) diff --git a/frontend/lib/src/render-tree/BlockNode.ts b/frontend/lib/src/render-tree/BlockNode.ts new file mode 100644 index 00000000000..69529ccd643 --- /dev/null +++ b/frontend/lib/src/render-tree/BlockNode.ts @@ -0,0 +1,196 @@ +/** + * Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2025) + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { Block as BlockProto, Element } from "@streamlit/protobuf" + +import { isNullOrUndefined, notUndefined } from "~lib/util/utils" + +import { AppNode, NO_SCRIPT_RUN_ID } from "./AppNode.interface" + +/** + * A container AppNode that holds children. + */ +export class BlockNode implements AppNode { + public readonly children: AppNode[] + + public readonly deltaBlock: BlockProto + + public readonly scriptRunId: string + + public readonly fragmentId?: string + + public readonly deltaMsgReceivedAt?: number + + // The hash of the script that created this block. + public readonly activeScriptHash: string + + public constructor( + activeScriptHash: string, + children?: AppNode[], + deltaBlock?: BlockProto, + scriptRunId?: string, + fragmentId?: string, + deltaMsgReceivedAt?: number + ) { + this.activeScriptHash = activeScriptHash + this.children = children ?? [] + this.deltaBlock = deltaBlock ?? new BlockProto({}) + this.scriptRunId = scriptRunId ?? NO_SCRIPT_RUN_ID + this.fragmentId = fragmentId + this.deltaMsgReceivedAt = deltaMsgReceivedAt + } + + /** True if this Block has no children. */ + public get isEmpty(): boolean { + return this.children.length === 0 + } + + public getIn(path: number[]): AppNode | undefined { + if (path.length === 0) { + return undefined + } + + const childIndex = path[0] + if (childIndex < 0 || childIndex >= this.children.length) { + return undefined + } + + if (path.length === 1) { + return this.children[childIndex] + } + + return this.children[childIndex].getIn(path.slice(1)) + } + + public setIn(path: number[], node: AppNode, scriptRunId: string): BlockNode { + if (path.length === 0) { + throw new Error(`empty path!`) + } + + const childIndex = path[0] + if (childIndex < 0 || childIndex > this.children.length) { + throw new Error( + `Bad 'setIn' index ${childIndex} (should be between [0, ${this.children.length}])` + ) + } + + const newChildren = this.children.slice() + if (path.length === 1) { + // Base case + newChildren[childIndex] = node + } else { + // Pop the current element off our path, and recurse into our children + newChildren[childIndex] = newChildren[childIndex].setIn( + path.slice(1), + node, + scriptRunId + ) + } + + return new BlockNode( + this.activeScriptHash, + newChildren, + this.deltaBlock, + scriptRunId, + this.fragmentId, + this.deltaMsgReceivedAt + ) + } + + filterMainScriptElements(mainScriptHash: string): AppNode | undefined { + if (this.activeScriptHash !== mainScriptHash) { + return undefined + } + + // Recursively clear our children. + const newChildren = this.children + .map(child => child.filterMainScriptElements(mainScriptHash)) + .filter(notUndefined) + + return new BlockNode( + this.activeScriptHash, + newChildren, + this.deltaBlock, + this.scriptRunId, + this.fragmentId, + this.deltaMsgReceivedAt + ) + } + + public clearStaleNodes( + currentScriptRunId: string, + fragmentIdsThisRun?: Array, + fragmentIdOfBlock?: string + ): BlockNode | undefined { + if (!fragmentIdsThisRun?.length) { + // If we're not currently running a fragment, then we can remove any blocks + // that don't correspond to currentScriptRunId. + if (this.scriptRunId !== currentScriptRunId) { + return undefined + } + } else { + // Otherwise, we are currently running a fragment, and our behavior + // depends on the fragmentId of this BlockNode. + + // The parent block was modified but this element wasn't, so it's stale. + if (fragmentIdOfBlock && this.scriptRunId !== currentScriptRunId) { + return undefined + } + + // This block is modified by the current run, so we indicate this to our children in case + // they were not modified by the current run, which means they are stale. + if ( + this.fragmentId && + fragmentIdsThisRun.includes(this.fragmentId) && + this.scriptRunId === currentScriptRunId + ) { + fragmentIdOfBlock = this.fragmentId + } + } + + // Recursively clear our children. + const newChildren = this.children + .map(child => { + return child.clearStaleNodes( + currentScriptRunId, + fragmentIdsThisRun, + fragmentIdOfBlock + ) + }) + .filter(notUndefined) + + return new BlockNode( + this.activeScriptHash, + newChildren, + this.deltaBlock, + currentScriptRunId, + this.fragmentId, + this.deltaMsgReceivedAt + ) + } + + public getElements(elementSet?: Set): Set { + if (isNullOrUndefined(elementSet)) { + elementSet = new Set() + } + + for (const child of this.children) { + child.getElements(elementSet) + } + + return elementSet + } +} diff --git a/frontend/lib/src/render-tree/ElementNode.test.ts b/frontend/lib/src/render-tree/ElementNode.test.ts new file mode 100644 index 00000000000..fa1ad904abe --- /dev/null +++ b/frontend/lib/src/render-tree/ElementNode.test.ts @@ -0,0 +1,421 @@ +/** + * Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2025) + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { ArrowNamedDataSet, IArrowVegaLiteChart } from "@streamlit/protobuf" + +import { UNICODE } from "~lib/mocks/arrow" + +import { NO_SCRIPT_RUN_ID } from "./AppNode.interface" +import { + arrowDataFrame, + arrowTable, + arrowVegaLiteChart, + text, +} from "./test-utils" + +describe("ElementNode", () => { + describe("ElementNode.quiverElement", () => { + it("returns a quiverElement (arrowTable)", () => { + const node = arrowTable() + const q = node.quiverElement + expect(q.columnNames).toEqual([["", "c1", "c2"]]) + expect(q.getCell(0, 0).content).toEqual("i1") + }) + + it("returns a quiverElement (arrowDataFrame)", () => { + const node = arrowDataFrame() + const q = node.quiverElement + expect(q.columnNames).toEqual([["", "c1", "c2"]]) + expect(q.getCell(0, 0).content).toEqual("i1") + }) + + it("does not recompute its value (arrowTable)", () => { + // accessing `quiverElement` twice should return the same instance. + const node = arrowTable() + expect(node.quiverElement).toStrictEqual(node.quiverElement) + }) + + it("does not recompute its value (arrowDataFrame)", () => { + // accessing `quiverElement` twice should return the same instance. + const node = arrowDataFrame() + expect(node.quiverElement).toStrictEqual(node.quiverElement) + }) + + it("throws an error for other element types", () => { + const node = text("foo") + expect(() => node.quiverElement).toThrow( + "elementType 'text' is not a valid Quiver element!" + ) + }) + }) + + describe("ElementNode.vegaLiteChartElement", () => { + it("returns a vegaLiteChartElement (data)", () => { + const MOCK_VEGA_LITE_CHART = { + spec: JSON.stringify({ + mark: "circle", + encoding: { + x: { field: "a", type: "quantitative" }, + y: { field: "b", type: "quantitative" }, + size: { field: "c", type: "quantitative" }, + color: { field: "c", type: "quantitative" }, + }, + }), + data: { data: UNICODE }, + datasets: [], + useContainerWidth: true, + } + const node = arrowVegaLiteChart(MOCK_VEGA_LITE_CHART) + const element = node.vegaLiteChartElement + + // spec + expect(element.spec).toEqual(MOCK_VEGA_LITE_CHART.spec) + + // data + expect(element.data?.columnNames).toEqual([["", "c1", "c2"]]) + expect(element.data?.getCell(0, 0).content).toEqual("i1") + + // datasets + expect(element.datasets.length).toEqual(0) + + // use container width + expect(element.useContainerWidth).toEqual( + MOCK_VEGA_LITE_CHART.useContainerWidth + ) + }) + + it("returns a vegaLiteChartElement (datasets)", () => { + const MOCK_VEGA_LITE_CHART = { + spec: JSON.stringify({ + mark: "circle", + encoding: { + x: { field: "a", type: "quantitative" }, + y: { field: "b", type: "quantitative" }, + size: { field: "c", type: "quantitative" }, + color: { field: "c", type: "quantitative" }, + }, + }), + data: null, + datasets: [{ hasName: true, name: "foo", data: { data: UNICODE } }], + useContainerWidth: true, + } + const node = arrowVegaLiteChart(MOCK_VEGA_LITE_CHART) + const element = node.vegaLiteChartElement + + // spec + expect(element.spec).toEqual(MOCK_VEGA_LITE_CHART.spec) + + // data + expect(element.data).toEqual(null) + + // datasets + expect(element.datasets[0].hasName).toEqual( + MOCK_VEGA_LITE_CHART.datasets[0].hasName + ) + expect(element.datasets[0].name).toEqual( + MOCK_VEGA_LITE_CHART.datasets[0].name + ) + expect(element.datasets[0].data.columnNames).toEqual([["", "c1", "c2"]]) + expect(element.datasets[0].data.getCell(0, 0).content).toEqual("i1") + + // use container width + expect(element.useContainerWidth).toEqual( + MOCK_VEGA_LITE_CHART.useContainerWidth + ) + }) + + it("does not recompute its value", () => { + const MOCK_VEGA_LITE_CHART = { + spec: JSON.stringify({ + mark: "circle", + encoding: { + x: { field: "a", type: "quantitative" }, + y: { field: "b", type: "quantitative" }, + size: { field: "c", type: "quantitative" }, + color: { field: "c", type: "quantitative" }, + }, + }), + data: { data: UNICODE }, + datasets: [], + useContainerWidth: true, + } + // accessing `vegaLiteChartElement` twice should return the same instance. + const node = arrowVegaLiteChart(MOCK_VEGA_LITE_CHART) + expect(node.vegaLiteChartElement).toStrictEqual( + node.vegaLiteChartElement + ) + }) + + it("throws an error for other element types", () => { + const node = text("foo") + expect(() => node.vegaLiteChartElement).toThrow( + "elementType 'text' is not a valid VegaLiteChartElement!" + ) + }) + }) + + describe("ElementNode.arrowAddRows", () => { + const MOCK_UNNAMED_DATASET = { + hasName: false, + name: "", + data: { data: UNICODE }, + } as ArrowNamedDataSet + const MOCK_NAMED_DATASET = { + hasName: true, + name: "foo", + data: { data: UNICODE }, + } as ArrowNamedDataSet + const MOCK_ANOTHER_NAMED_DATASET = { + hasName: true, + name: "bar", + data: { data: UNICODE }, + } as ArrowNamedDataSet + + describe("arrowTable", () => { + test("addRows can be called with an unnamed dataset", () => { + const node = arrowTable() + const newNode = node.arrowAddRows( + MOCK_UNNAMED_DATASET, + NO_SCRIPT_RUN_ID + ) + const q = newNode.quiverElement + + expect(q.columnNames).toEqual([["", "c1", "c2"]]) + expect(q.dimensions.numDataRows).toEqual(4) + expect(q.getCell(0, 0).content).toEqual("i1") + expect(q.getCell(2, 0).content).toEqual("i1") + expect(q.getCell(0, 1).content).toEqual("foo") + expect(q.getCell(2, 1).content).toEqual("foo") + }) + + test("addRows throws an error when called with a named dataset", () => { + const node = arrowTable() + expect(() => + node.arrowAddRows(MOCK_NAMED_DATASET, NO_SCRIPT_RUN_ID) + ).toThrow( + "Add rows cannot be used with a named dataset for this element." + ) + }) + }) + + describe("arrowDataFrame", () => { + test("addRows can be called with an unnamed dataset", () => { + const node = arrowDataFrame() + const newNode = node.arrowAddRows( + MOCK_UNNAMED_DATASET, + NO_SCRIPT_RUN_ID + ) + const q = newNode.quiverElement + + expect(q.columnNames).toEqual([["", "c1", "c2"]]) + expect(q.dimensions.numDataRows).toEqual(4) + expect(q.getCell(0, 0).content).toEqual("i1") + expect(q.getCell(2, 0).content).toEqual("i1") + expect(q.getCell(0, 1).content).toEqual("foo") + expect(q.getCell(2, 1).content).toEqual("foo") + }) + + test("addRows throws an error when called with a named dataset", () => { + const node = arrowDataFrame() + expect(() => + node.arrowAddRows(MOCK_NAMED_DATASET, NO_SCRIPT_RUN_ID) + ).toThrow( + "Add rows cannot be used with a named dataset for this element." + ) + }) + }) + + describe("arrowVegaLiteChart", () => { + const getVegaLiteChart = ( + datasets?: ArrowNamedDataSet[], + data?: Uint8Array + ): IArrowVegaLiteChart => ({ + datasets: datasets || [], + data: data ? { data } : null, + spec: JSON.stringify({ + mark: "circle", + encoding: { + x: { field: "a", type: "quantitative" }, + y: { field: "b", type: "quantitative" }, + size: { field: "c", type: "quantitative" }, + color: { field: "c", type: "quantitative" }, + }, + }), + useContainerWidth: true, + }) + + describe("addRows is called with a named dataset", () => { + test("element has one dataset -> append new rows to that dataset", () => { + const node = arrowVegaLiteChart( + getVegaLiteChart([MOCK_ANOTHER_NAMED_DATASET]) + ) + const newNode = node.arrowAddRows( + MOCK_NAMED_DATASET, + NO_SCRIPT_RUN_ID + ) + const element = newNode.vegaLiteChartElement + + const quiverData = element.datasets[0].data + expect(quiverData?.columnNames).toEqual([["", "c1", "c2"]]) + expect(quiverData?.dimensions.numDataRows).toEqual(4) + + expect(quiverData?.getCell(0, 0).content).toEqual("i1") + expect(quiverData?.getCell(0, 1).content).toEqual("foo") + expect(quiverData?.getCell(2, 0).content).toEqual("i1") + expect(quiverData?.getCell(2, 1).content).toEqual("foo") + }) + + test("element has a dataset with the given name -> append new rows to that dataset", () => { + const node = arrowVegaLiteChart( + getVegaLiteChart([MOCK_NAMED_DATASET, MOCK_ANOTHER_NAMED_DATASET]) + ) + const newNode = node.arrowAddRows( + MOCK_NAMED_DATASET, + NO_SCRIPT_RUN_ID + ) + const element = newNode.vegaLiteChartElement + + const quiverData = element.datasets[0].data + expect(quiverData?.columnNames).toEqual([["", "c1", "c2"]]) + expect(quiverData?.dimensions.numDataRows).toEqual(4) + + expect(quiverData?.getCell(0, 0).content).toEqual("i1") + expect(quiverData?.getCell(0, 1).content).toEqual("foo") + expect(quiverData?.getCell(2, 0).content).toEqual("i1") + expect(quiverData?.getCell(2, 1).content).toEqual("foo") + }) + + test("element doesn't have a matched dataset, but has data -> append new rows to data", () => { + const node = arrowVegaLiteChart(getVegaLiteChart(undefined, UNICODE)) + const newNode = node.arrowAddRows( + MOCK_NAMED_DATASET, + NO_SCRIPT_RUN_ID + ) + const element = newNode.vegaLiteChartElement + + const quiverData = element.data + expect(quiverData?.columnNames).toEqual([["", "c1", "c2"]]) + expect(quiverData?.dimensions.numDataRows).toEqual(4) + + expect(quiverData?.getCell(0, 0).content).toEqual("i1") + expect(quiverData?.getCell(0, 1).content).toEqual("foo") + expect(quiverData?.getCell(2, 0).content).toEqual("i1") + expect(quiverData?.getCell(2, 1).content).toEqual("foo") + }) + + test("element doesn't have a matched dataset or data -> use new rows as data", () => { + const node = arrowVegaLiteChart( + getVegaLiteChart([ + MOCK_ANOTHER_NAMED_DATASET, + MOCK_ANOTHER_NAMED_DATASET, + ]) + ) + const newNode = node.arrowAddRows( + MOCK_NAMED_DATASET, + NO_SCRIPT_RUN_ID + ) + const element = newNode.vegaLiteChartElement + + const quiverData = element.data + expect(quiverData?.columnNames).toEqual([["", "c1", "c2"]]) + expect(quiverData?.dimensions.numDataRows).toEqual(2) + + expect(quiverData?.getCell(0, 0).content).toEqual("i1") + expect(quiverData?.getCell(0, 1).content).toEqual("foo") + }) + + test("element doesn't have any datasets or data -> use new rows as data", () => { + const node = arrowVegaLiteChart(getVegaLiteChart()) + const newNode = node.arrowAddRows( + MOCK_NAMED_DATASET, + NO_SCRIPT_RUN_ID + ) + const element = newNode.vegaLiteChartElement + + const quiverData = element.data + expect(quiverData?.columnNames).toEqual([["", "c1", "c2"]]) + expect(quiverData?.dimensions.numDataRows).toEqual(2) + + expect(quiverData?.getCell(0, 0).content).toEqual("i1") + expect(quiverData?.getCell(0, 1).content).toEqual("foo") + }) + }) + + describe("addRows is called with an unnamed dataset", () => { + test("element has one dataset -> append new rows to that dataset", () => { + const node = arrowVegaLiteChart( + getVegaLiteChart([MOCK_NAMED_DATASET]) + ) + const newNode = node.arrowAddRows( + MOCK_UNNAMED_DATASET, + NO_SCRIPT_RUN_ID + ) + const element = newNode.vegaLiteChartElement + + const quiverData = element.datasets[0].data + expect(quiverData.columnNames).toEqual([["", "c1", "c2"]]) + expect(quiverData?.dimensions.numDataRows).toEqual(4) + + expect(quiverData.getCell(0, 0).content).toEqual("i1") + expect(quiverData.getCell(2, 0).content).toEqual("i1") + expect(quiverData.getCell(0, 1).content).toEqual("foo") + expect(quiverData.getCell(2, 1).content).toEqual("foo") + }) + + test("element has data -> append new rows to data", () => { + const node = arrowVegaLiteChart(getVegaLiteChart(undefined, UNICODE)) + const newNode = node.arrowAddRows( + MOCK_UNNAMED_DATASET, + NO_SCRIPT_RUN_ID + ) + const element = newNode.vegaLiteChartElement + + const quiverData = element.data + expect(quiverData?.columnNames).toEqual([["", "c1", "c2"]]) + expect(quiverData?.dimensions.numDataRows).toEqual(4) + + expect(quiverData?.getCell(0, 0).content).toEqual("i1") + expect(quiverData?.getCell(2, 0).content).toEqual("i1") + expect(quiverData?.getCell(0, 1).content).toEqual("foo") + expect(quiverData?.getCell(2, 1).content).toEqual("foo") + }) + + test("element doesn't have any datasets or data -> use new rows as data", () => { + const node = arrowVegaLiteChart(getVegaLiteChart()) + const newNode = node.arrowAddRows( + MOCK_UNNAMED_DATASET, + NO_SCRIPT_RUN_ID + ) + const element = newNode.vegaLiteChartElement + + const quiverData = element.data + expect(quiverData?.columnNames).toEqual([["", "c1", "c2"]]) + expect(quiverData?.dimensions.numDataRows).toEqual(2) + + expect(quiverData?.getCell(0, 0).content).toEqual("i1") + expect(quiverData?.getCell(0, 1).content).toEqual("foo") + }) + }) + }) + + it("throws an error for other element types", () => { + const node = text("foo") + expect(() => + node.arrowAddRows(MOCK_UNNAMED_DATASET, NO_SCRIPT_RUN_ID) + ).toThrow("elementType 'text' is not a valid arrowAddRows target!") + }) + }) +}) diff --git a/frontend/lib/src/render-tree/ElementNode.ts b/frontend/lib/src/render-tree/ElementNode.ts new file mode 100644 index 00000000000..958fdf7af17 --- /dev/null +++ b/frontend/lib/src/render-tree/ElementNode.ts @@ -0,0 +1,273 @@ +/** + * Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2025) + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { produce } from "immer" + +import { + ArrowNamedDataSet, + Arrow as ArrowProto, + ArrowVegaLiteChart as ArrowVegaLiteChartProto, + Element, + ForwardMsgMetadata, + IArrow, + IArrowNamedDataSet, +} from "@streamlit/protobuf" + +import { + VegaLiteChartElement, + WrappedNamedDataset, +} from "~lib/components/elements/ArrowVegaLiteChart" +import { Quiver } from "~lib/dataframes/Quiver" +import { isNullOrUndefined } from "~lib/util/utils" + +import { AppNode } from "./AppNode.interface" + +/** + * A leaf AppNode. Contains a single element to render. + */ +export class ElementNode implements AppNode { + public readonly element: Element + + public readonly metadata: ForwardMsgMetadata + + public readonly scriptRunId: string + + public readonly fragmentId?: string + + private lazyQuiverElement?: Quiver + + private lazyVegaLiteChartElement?: VegaLiteChartElement + + // The hash of the script that created this element. + public readonly activeScriptHash: string + + /** Create a new ElementNode. */ + public constructor( + element: Element, + metadata: ForwardMsgMetadata, + scriptRunId: string, + activeScriptHash: string, + fragmentId?: string + ) { + this.element = element + this.metadata = metadata + this.scriptRunId = scriptRunId + this.activeScriptHash = activeScriptHash + this.fragmentId = fragmentId + } + + public get quiverElement(): Quiver { + if (this.lazyQuiverElement !== undefined) { + return this.lazyQuiverElement + } + + if ( + this.element.type !== "arrowTable" && + this.element.type !== "arrowDataFrame" + ) { + throw new Error( + `elementType '${this.element.type}' is not a valid Quiver element!` + ) + } + + const toReturn = new Quiver(this.element[this.element.type] as ArrowProto) + // TODO (lukasmasuch): Delete element from proto object? + this.lazyQuiverElement = toReturn + return toReturn + } + + public get vegaLiteChartElement(): VegaLiteChartElement { + if (this.lazyVegaLiteChartElement !== undefined) { + return this.lazyVegaLiteChartElement + } + + if (this.element.type !== "arrowVegaLiteChart") { + throw new Error( + `elementType '${this.element.type}' is not a valid VegaLiteChartElement!` + ) + } + + const proto = this.element.arrowVegaLiteChart as ArrowVegaLiteChartProto + const modifiedData = proto.data ? new Quiver(proto.data) : null + const modifiedDatasets = + proto.datasets.length > 0 ? wrapDatasets(proto.datasets) : [] + + const toReturn = { + data: modifiedData, + spec: proto.spec, + datasets: modifiedDatasets, + useContainerWidth: proto.useContainerWidth, + vegaLiteTheme: proto.theme, + id: proto.id, + selectionMode: proto.selectionMode, + formId: proto.formId, + } + + this.lazyVegaLiteChartElement = toReturn + return toReturn + } + + public getIn(): AppNode | undefined { + return undefined + } + + public setIn(): AppNode { + throw new Error("'setIn' cannot be called on an ElementNode") + } + + public filterMainScriptElements( + mainScriptHash: string + ): AppNode | undefined { + if (this.activeScriptHash !== mainScriptHash) { + return undefined + } + + return this + } + + public clearStaleNodes( + currentScriptRunId: string, + fragmentIdsThisRun?: Array, + fragmentIdOfBlock?: string + ): ElementNode | undefined { + if (fragmentIdsThisRun?.length) { + // If we're currently running a fragment, nodes unrelated to the fragment + // shouldn't be cleared. This can happen when, + // 1. This element doesn't correspond to a fragment at all. + // 2. This element is a fragment but is in no path that was modified. + // 3. This element belongs to a path that was modified, but it was modified in the same run. + if ( + !this.fragmentId || + !fragmentIdOfBlock || + this.scriptRunId === currentScriptRunId + ) { + return this + } + } + return this.scriptRunId === currentScriptRunId ? this : undefined + } + + public getElements(elements?: Set): Set { + if (isNullOrUndefined(elements)) { + elements = new Set() + } + elements.add(this.element) + return elements + } + + public arrowAddRows( + namedDataSet: ArrowNamedDataSet, + scriptRunId: string + ): ElementNode { + const elementType = this.element.type + const newNode = new ElementNode( + this.element, + this.metadata, + scriptRunId, + this.activeScriptHash, + this.fragmentId + ) + + switch (elementType) { + case "arrowTable": + case "arrowDataFrame": { + newNode.lazyQuiverElement = ElementNode.quiverAddRowsHelper( + this.quiverElement, + namedDataSet + ) + break + } + case "arrowVegaLiteChart": { + newNode.lazyVegaLiteChartElement = + ElementNode.vegaLiteChartAddRowsHelper( + this.vegaLiteChartElement, + namedDataSet + ) + break + } + default: { + // This should never happen! + throw new Error( + `elementType '${this.element.type}' is not a valid arrowAddRows target!` + ) + } + } + + return newNode + } + + private static quiverAddRowsHelper( + element: Quiver, + namedDataSet: ArrowNamedDataSet + ): Quiver { + if (namedDataSet.hasName) { + throw new Error( + "Add rows cannot be used with a named dataset for this element." + ) + } + + const newQuiver = new Quiver(namedDataSet.data as IArrow) + return element.addRows(newQuiver) + } + + private static vegaLiteChartAddRowsHelper( + element: VegaLiteChartElement, + namedDataSet: ArrowNamedDataSet + ): VegaLiteChartElement { + const newDataSetName = namedDataSet.hasName ? namedDataSet.name : null + const newDataSetQuiver = new Quiver(namedDataSet.data as IArrow) + + return produce(element, (draft: VegaLiteChartElement) => { + const existingDataSet = getNamedDataSet(draft.datasets, newDataSetName) + if (existingDataSet) { + existingDataSet.data = existingDataSet.data.addRows(newDataSetQuiver) + } else { + draft.data = draft.data + ? draft.data.addRows(newDataSetQuiver) + : newDataSetQuiver + } + }) + } +} + +/** + * If there is only one NamedDataSet, return it. + * If there is a NamedDataset that matches the given name, return it. + * Otherwise, return `undefined`. + */ +function getNamedDataSet( + namedDataSets: WrappedNamedDataset[], + name: string | null +): WrappedNamedDataset | undefined { + if (namedDataSets.length === 1) { + return namedDataSets[0] + } + + return namedDataSets.find( + (dataset: WrappedNamedDataset) => dataset.hasName && dataset.name === name + ) +} + +/** Iterates over datasets and converts data to Quiver. */ +function wrapDatasets(datasets: IArrowNamedDataSet[]): WrappedNamedDataset[] { + return datasets.map((dataset: IArrowNamedDataSet) => { + return { + hasName: dataset.hasName as boolean, + name: dataset.name as string, + data: new Quiver(dataset.data as IArrow), + } + }) +} diff --git a/frontend/lib/src/render-tree/test-utils.ts b/frontend/lib/src/render-tree/test-utils.ts new file mode 100644 index 00000000000..bf245c630b8 --- /dev/null +++ b/frontend/lib/src/render-tree/test-utils.ts @@ -0,0 +1,173 @@ +/** + * Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2025) + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { Writer } from "protobufjs" + +import { + Block as BlockProto, + Element, + ForwardMsgMetadata, + IArrowVegaLiteChart, +} from "@streamlit/protobuf" + +import { UNICODE } from "~lib/mocks/arrow" +import { isNullOrUndefined } from "~lib/util/utils" + +import { AppNode, NO_SCRIPT_RUN_ID } from "./AppNode.interface" +import { BlockNode } from "./BlockNode" +import { ElementNode } from "./ElementNode" + +export const FAKE_SCRIPT_HASH = "fake_script_hash" + +/** Create a `Text` element node with the given properties. */ +export function text( + textArg: string, + scriptRunId = NO_SCRIPT_RUN_ID +): ElementNode { + const element = makeProto(Element, { text: { body: textArg } }) + return new ElementNode( + element, + ForwardMsgMetadata.create(), + scriptRunId, + FAKE_SCRIPT_HASH + ) +} + +/** Create a BlockNode with the given properties. */ +export function block( + children: AppNode[] = [], + scriptRunId = NO_SCRIPT_RUN_ID +): BlockNode { + return new BlockNode( + FAKE_SCRIPT_HASH, + children, + makeProto(BlockProto, {}), + scriptRunId + ) +} + +/** Create an arrowTable element node with the given properties. */ +export function arrowTable(scriptRunId = NO_SCRIPT_RUN_ID): ElementNode { + const element = makeProto(Element, { arrowTable: { data: UNICODE } }) + return new ElementNode( + element, + ForwardMsgMetadata.create(), + scriptRunId, + FAKE_SCRIPT_HASH + ) +} + +/** Create an arrowDataFrame element node with the given properties. */ +export function arrowDataFrame(scriptRunId = NO_SCRIPT_RUN_ID): ElementNode { + const element = makeProto(Element, { arrowDataFrame: { data: UNICODE } }) + return new ElementNode( + element, + ForwardMsgMetadata.create(), + scriptRunId, + FAKE_SCRIPT_HASH + ) +} + +/** Create an arrowVegaLiteChart element node with the given properties. */ +export function arrowVegaLiteChart( + data: IArrowVegaLiteChart, + scriptRunId = NO_SCRIPT_RUN_ID +): ElementNode { + const element = makeProto(Element, { arrowVegaLiteChart: data }) + return new ElementNode( + element, + ForwardMsgMetadata.create(), + scriptRunId, + FAKE_SCRIPT_HASH + ) +} + +/** Create a ForwardMsgMetadata with the given container and path */ +export function forwardMsgMetadata( + deltaPath: number[], + activeScriptHash = FAKE_SCRIPT_HASH +): ForwardMsgMetadata { + expect(deltaPath.length).toBeGreaterThanOrEqual(2) + return makeProto(ForwardMsgMetadata, { deltaPath, activeScriptHash }) +} + +/** + * Make a "fully concrete" instance of a protobuf message. + * This function constructs a message and then encodes and decodes it as + * if it had arrived on the wire. This ensures that that it has all its + * 'oneOfs' and 'defaults' set. + */ +export function makeProto( + MessageType: { + new (props: Props): Type + encode: (message: Type, writer: Writer) => Writer + decode: (bytes: Uint8Array) => Type + }, + properties: Props +): Type { + const message = new MessageType(properties) + const bytes = MessageType.encode(message, Writer.create()).finish() + return MessageType.decode(bytes) +} + +// Custom Jest matchers for dealing with AppNodes +declare global { + // eslint-disable-next-line @typescript-eslint/no-namespace + namespace vi { + interface Matchers { + toBeTextNode(text: string): R + } + } +} + +interface CustomMatchers { + toBeTextNode(text: string): R +} + +declare module "vitest" { + // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-empty-object-type -- TODO: Replace 'any' with a more specific type. + interface Assertion extends CustomMatchers {} + // eslint-disable-next-line @typescript-eslint/no-empty-object-type + interface AsymmetricMatchersContaining extends CustomMatchers {} +} + +expect.extend({ + toBeTextNode(received, textArg) { + const elementNode = received as ElementNode + if (isNullOrUndefined(elementNode)) { + return { + message: () => `expected ${received} to be an instance of ElementNode`, + pass: false, + } + } + + const { type } = elementNode.element + if (type !== "text") { + return { + message: () => + `expected ${received}.element.type to be 'text', but it was ${type}`, + pass: false, + } + } + + const textBody = elementNode.element.text?.body + return { + message: () => + `expected ${received}.element.text.body to be "${textArg}", but it was "${textBody}"`, + pass: textBody === textArg, + } + }, +}) diff --git a/frontend/vitest.config.ts b/frontend/vitest.config.ts index 7508481cf38..e920f77aa8b 100644 --- a/frontend/vitest.config.ts +++ b/frontend/vitest.config.ts @@ -26,7 +26,11 @@ export default defineConfig({ provider: "v8", reporter: ["text-summary", "json-summary", "html"], include: ["*/src/**/*"], - exclude: ["lib/src/vendor/**", ...coverageConfigDefaults.exclude], + exclude: [ + "lib/src/vendor/**", + "**/*.interface.ts", + ...coverageConfigDefaults.exclude, + ], }, }, })