Cache JSON when using pydeck local data (streamlit#7113)

## Describe your changes - add a protobuf key named element id which is a hash of the pydeck json - add logic to not parse the json again where the json is stored in state - if the hash string is different, then parse it again - if it’s not, use the cached json - add fullScreen and theme to state - if these change, we need to parse the json again to rerender the new map - these changes make the loading slightly longer when rerunning or running the script; however, the experience is significantly better so it's worth to do this. ## GitHub Issue Link (if applicable) streamlit#5532 ## Testing Plan - Explanation of why no additional tests are needed - Unit Tests (JS and/or Python) - will need to add these - E2E Tests - very difficult to add these so will not be adding these. - Any manual testing needed? - yes, i manually tested the test case within the issue 5532. I will grab other use cases from pydeck and make sure they're responsive. This branch: https://github.com/streamlit/streamlit/assets/16749069/130d21ca-efe4-437b-aa6e-fb1a04c6805b Develop: https://github.com/streamlit/streamlit/assets/16749069/275b4e00-4f1d-444e-bf69-f84040c9b86f --- **Contribution License Agreement** By submitting this pull request you agree that all contributions to this project are made under the Apache 2.0 license.
zyxue · Apr 16, 2024 · b12f855 · b12f855
1 parent da884c6
commit b12f855
Show file tree

Hide file tree

Showing 5 changed files with 191 additions and 84 deletions.
diff --git a/frontend/lib/src/components/elements/DeckGlJsonChart/DeckGlJsonChart.test.tsx b/frontend/lib/src/components/elements/DeckGlJsonChart/DeckGlJsonChart.test.tsx
@@ -15,28 +15,36 @@
  */
 
 import React from "react"
-import { DeckGL } from "deck.gl"
-import { shallow } from "@streamlit/lib/src/test_util"
+import { render } from "@streamlit/lib/src/test_util"
 
 import { DeckGlJsonChart as DeckGlJsonChartProto } from "@streamlit/lib/src/proto"
-import { NavigationControl } from "react-map-gl"
+import { screen } from "@testing-library/react"
+import "@testing-library/jest-dom"
 import { mockTheme } from "@streamlit/lib/src/mocks/mockTheme"
-import { DeckGlJsonChart, PropsWithHeight } from "./DeckGlJsonChart"
+import { DeckGlJsonChart, PropsWithHeight, State } from "./DeckGlJsonChart"
+
+const mockInitialViewState = {
+  bearing: -27.36,
+  latitude: 52.2323,
+  longitude: -1.415,
+  maxZoom: 15,
+  minZoom: 5,
+  pitch: 40.5,
+  height: 500,
+  zoom: 6,
+}
+
+const mockId = "testId"
+jest.mock("@streamlit/lib/src/theme", () => ({
+  hasLightBackgroundColor: jest.fn(() => false),
+}))
 
 const getProps = (
   elementProps: Partial<DeckGlJsonChartProto> = {},
   initialViewStateProps: Record<string, unknown> = {}
 ): PropsWithHeight => {
   const json = {
-    initialViewState: {
-      bearing: -27.36,
-      latitude: 52.2323,
-      longitude: -1.415,
-      maxZoom: 15,
-      minZoom: 5,
-      pitch: 40.5,
-      zoom: 6,
-    },
+    initialViewState: mockInitialViewState,
     layers: [
       {
         "@@type": "HexagonLayer",
@@ -62,84 +70,151 @@ const getProps = (
 
   return {
     element: DeckGlJsonChartProto.create({
+      id: mockId,
       json: JSON.stringify(json),
       ...elementProps,
     }),
     width: 0,
     mapboxToken: "mapboxToken",
     height: undefined,
     theme: mockTheme.emotion,
+    isFullScreen: false,
   }
 }
 
 describe("DeckGlJsonChart element", () => {
   it("renders without crashing", () => {
     const props = getProps()
-    const wrapper = shallow(<DeckGlJsonChart {...props} />)
 
-    expect(wrapper.find(DeckGL).length).toBe(1)
-    expect(wrapper.find(NavigationControl).length).toBe(1)
+    render(<DeckGlJsonChart {...props} />)
+
+    const deckGlJsonChart = screen.getByTestId("stDeckGlJsonChart")
+    expect(deckGlJsonChart).toBeInTheDocument()
   })
 
-  it("merges client and server changes", () => {
+  it("should merge client and server changes in viewState", () => {
     const props = getProps()
-    const wrapper = shallow(<DeckGlJsonChart {...props} />)
-    const deckGL = wrapper.find(DeckGL)
+    const initialViewStateServer = mockInitialViewState
 
-    expect(deckGL.length).toBe(1)
-    expect(deckGL.prop("onViewStateChange")).toBeDefined()
+    const initialViewStateClient = { ...mockInitialViewState, zoom: 8 }
 
-    // @ts-expect-error
-    deckGL.prop("onViewStateChange")({
-      viewState: { pitch: 5, zoom: 5 },
-    })
+    const state = {
+      viewState: initialViewStateClient,
+      initialViewState: initialViewStateClient,
+    }
 
-    // @ts-expect-error
-    wrapper.setProps(getProps({}, { pitch: 40.5, zoom: 10 }))
+    const result = DeckGlJsonChart.getDerivedStateFromProps(props, state)
 
-    expect(wrapper.state("viewState")).toStrictEqual({
-      pitch: 5,
-      zoom: 10,
+    expect(result).toEqual({
+      // should match original mockInitialViewState
+      viewState: { ...initialViewStateClient, zoom: 6 },
+      initialViewState: initialViewStateServer,
     })
   })
 
-  it("should render tooltip", () => {
-    const props = getProps({
-      tooltip: `{"html": "<b>Elevation Value:</b> {elevationValue}", "style": {"color": "white"}}`,
+  describe("createTooltip", () => {
+    let deckGlInstance: any
+
+    beforeEach(() => {
+      deckGlInstance = new DeckGlJsonChart({ ...getProps() })
     })
-    const wrapper = shallow(<DeckGlJsonChart {...props} />)
-    const deckGL = wrapper.find(DeckGL)
 
-    expect(deckGL.length).toBe(1)
-    expect(deckGL.prop("getTooltip")).toBeDefined()
+    it("should return false if info is undefined", () => {
+      const result = deckGlInstance.createTooltip(undefined)
+      expect(result).toBe(false)
+    })
 
-    // @ts-expect-error
-    const createdTooltip = deckGL.prop("getTooltip")({
-      object: {
-        elevationValue: 10,
-      },
+    it("should return false if info.object is undefined", () => {
+      const result = deckGlInstance.createTooltip({})
+      expect(result).toBe(false)
     })
 
-    expect(createdTooltip.html).toBe("<b>Elevation Value:</b> 10")
-  })
+    it("should return false if element.tooltip is undefined", () => {
+      const result = deckGlInstance.createTooltip({ object: {} })
+      expect(result).toBe(false)
+    })
 
-  it("should render an empty tooltip", () => {
-    const props = getProps({
-      tooltip: "",
+    it("should interpolate the html with the correct object", () => {
+      deckGlInstance.props.element.tooltip = JSON.stringify({
+        html: "<b>Elevation Value:</b> {elevationValue}",
+      })
+      const result = deckGlInstance.createTooltip({
+        object: { elevationValue: 10 },
+      })
+
+      expect(result.html).toBe("<b>Elevation Value:</b> 10")
     })
-    const wrapper = shallow(<DeckGlJsonChart {...props} />)
-    const deckGL = wrapper.find(DeckGL)
 
-    expect(deckGL.length).toBe(1)
-    expect(deckGL.prop("getTooltip")).toBeDefined()
+    it("should interpolate the html with the an empty string", () => {
+      deckGlInstance.props = getProps({
+        tooltip: "",
+      })
+      const result = deckGlInstance.createTooltip({
+        object: { elevationValue: 10 },
+      })
 
-    // @ts-expect-error
-    const createdTooltip = deckGL.prop("getTooltip")({
-      object: {
-        elevationValue: 10,
-      },
+      expect(result.html).toBe(undefined)
+    })
+  })
+
+  describe("getDeckObject", () => {
+    const newId = "newTestId"
+    const newJson = {
+      initialViewState: mockInitialViewState,
+      mapStyle: "mapbox://styles/mapbox/light-v9",
+    }
+
+    const originalState: State = {
+      pydeckJson: newJson,
+      isFullScreen: false,
+      viewState: {},
+      initialized: false,
+      initialViewState: mockInitialViewState,
+      id: mockId,
+      isLightTheme: false,
+    }
+
+    const mockJsonParse = jest.fn().mockReturnValue(newJson)
+
+    beforeEach(() => {
+      JSON.parse = mockJsonParse
     })
 
-    expect(createdTooltip).toBe(false)
+    afterEach(() => {
+      mockJsonParse.mockClear()
+    })
+
+    const testJsonParsing = (
+      description: string,
+      stateOverride: Partial<State>
+    ): void => {
+      // the description will be passed in
+      // eslint-disable-next-line jest/valid-title
+      it(description, () => {
+        DeckGlJsonChart.getDeckObject(getProps(), originalState)
+
+        expect(JSON.parse).not.toHaveBeenCalled()
+
+        DeckGlJsonChart.getDeckObject(getProps(), {
+          ...originalState,
+          ...stateOverride,
+        })
+
+        expect(JSON.parse).toHaveBeenCalledTimes(1)
+      })
+    }
+
+    testJsonParsing(
+      "should call JSON.parse when the element id is different",
+      { id: newId }
+    )
+    testJsonParsing("should call JSON.parse when FullScreen state changes", {
+      id: mockId,
+      isFullScreen: true,
+    })
+    testJsonParsing("should call JSON.parse when theme state changes", {
+      id: mockId,
+      isLightTheme: true,
+    })
   })
 })
diff --git a/frontend/lib/src/components/elements/DeckGlJsonChart/DeckGlJsonChart.tsx b/frontend/lib/src/components/elements/DeckGlJsonChart/DeckGlJsonChart.tsx
@@ -39,8 +39,6 @@ import { registerLoaders } from "@loaders.gl/core"
 import withFullScreenWrapper from "@streamlit/lib/src/hocs/withFullScreenWrapper"
 import withMapboxToken from "@streamlit/lib/src/hocs/withMapboxToken"
 
-import { notNullOrUndefined } from "@streamlit/lib/src/util/utils"
-
 import { DeckGlJsonChart as DeckGlJsonChartProto } from "@streamlit/lib/src/proto"
 import {
   StyledDeckGlChart,
@@ -72,21 +70,26 @@ registerLoaders([CSVLoader, GLTFLoader])
 
 const jsonConverter = new JSONConverter({ configuration })
 
-interface Props {
+export interface DeckGLProps {
   width: number
   theme: EmotionTheme
   mapboxToken: string
   element: DeckGlJsonChartProto
+  isFullScreen?: boolean
 }
 
-export interface PropsWithHeight extends Props {
+export interface PropsWithHeight extends DeckGLProps {
   height?: number
 }
 
-interface State {
+export interface State {
   viewState: Record<string, unknown>
   initialized: boolean
   initialViewState: Record<string, unknown>
+  id: string | undefined
+  pydeckJson: any
+  isFullScreen: boolean
+  isLightTheme: boolean
 }
 
 export const DEFAULT_DECK_GL_HEIGHT = 500
@@ -100,6 +103,10 @@ export class DeckGlJsonChart extends PureComponent<PropsWithHeight, State> {
     },
     initialized: false,
     initialViewState: {},
+    id: undefined,
+    pydeckJson: undefined,
+    isFullScreen: false,
+    isLightTheme: hasLightBackgroundColor(this.props.theme),
   }
 
   componentDidMount = (): void => {
@@ -115,7 +122,7 @@ export class DeckGlJsonChart extends PureComponent<PropsWithHeight, State> {
     props: Readonly<PropsWithHeight>,
     state: Partial<State>
   ): Partial<State> | null {
-    const deck = DeckGlJsonChart.getDeckObject(props)
+    const deck = DeckGlJsonChart.getDeckObject(props, state)
 
     // If the ViewState on the server has changed, apply the diff to the current state
     if (!isEqual(deck.initialViewState, state.initialViewState)) {
@@ -144,35 +151,50 @@ export class DeckGlJsonChart extends PureComponent<PropsWithHeight, State> {
     return null
   }
 
-  static getDeckObject = (props: PropsWithHeight): DeckObject => {
-    const { element, width, height, theme } = props
-    const json = JSON.parse(element.json)
+  static getDeckObject = (
+    props: PropsWithHeight,
+    state: Partial<State>
+  ): DeckObject => {
+    const { element, width, height, theme, isFullScreen } = props
+
+    const currFullScreen = isFullScreen ?? false
+
+    // Only parse JSON when not transitioning to/from fullscreen, the element id changes, or theme changes
+    if (
+      element.id !== state.id ||
+      state.isFullScreen !== currFullScreen ||
+      state.isLightTheme !== hasLightBackgroundColor(theme)
+    ) {
+      state.pydeckJson = JSON.parse(element.json)
+      state.id = element.id
+    }
 
     // If unset, use either the Mapbox light or dark style based on Streamlit's theme
     // For Mapbox styles, see https://docs.mapbox.com/api/maps/styles/#mapbox-styles
-    if (!notNullOrUndefined(json.mapStyle)) {
-      const mapTheme = hasLightBackgroundColor(theme) ? "light" : "dark"
-      json.mapStyle = `mapbox://styles/mapbox/${mapTheme}-v9`
+    if (!state.pydeckJson?.mapStyle) {
+      state.pydeckJson.mapStyle = `mapbox://styles/mapbox/${
+        hasLightBackgroundColor(theme) ? "light" : "dark"
+      }-v9`
     }
 
-    // The graph dimensions could be set from props ( like withFullscreen ) or
-    // from the generated element object
-    if (height) {
-      json.initialViewState.height = height
-      json.initialViewState.width = width
+    // Set width and height based on the fullscreen state
+    if (isFullScreen) {
+      Object.assign(state.pydeckJson?.initialViewState, { width, height })
     } else {
-      if (!json.initialViewState.height) {
-        json.initialViewState.height = DEFAULT_DECK_GL_HEIGHT
+      if (!state.pydeckJson?.initialViewState?.height) {
+        state.pydeckJson.initialViewState.height = DEFAULT_DECK_GL_HEIGHT
       }
-
       if (element.useContainerWidth) {
-        json.initialViewState.width = width
+        state.pydeckJson.initialViewState.width = width
       }
     }
 
-    delete json.views // We are not using views. This avoids a console warning.
+    state.isFullScreen = isFullScreen
+    state.isLightTheme = hasLightBackgroundColor(theme)
+
+    delete state.pydeckJson?.views // We are not using views. This avoids a console warning.
 
-    return jsonConverter.convert(json)
+    return jsonConverter.convert(state.pydeckJson)
   }
 
   createTooltip = (info: PickingInfo): Record<string, unknown> | boolean => {
@@ -213,14 +235,14 @@ export class DeckGlJsonChart extends PureComponent<PropsWithHeight, State> {
   }
 
   render(): ReactNode {
-    const deck = DeckGlJsonChart.getDeckObject(this.props)
+    const deck = DeckGlJsonChart.getDeckObject(this.props, this.state)
     const { viewState } = this.state
-
     return (
       <StyledDeckGlChart
         className="stDeckGlJsonChart"
         width={deck.initialViewState.width}
         height={deck.initialViewState.height}
+        data-testid="stDeckGlJsonChart"
       >
         <DeckGL
           viewState={viewState}

diff --git a/frontend/lib/src/hocs/withFullScreenWrapper/withFullScreenWrapper.tsx b/frontend/lib/src/hocs/withFullScreenWrapper/withFullScreenWrapper.tsx
@@ -19,7 +19,7 @@ import hoistNonReactStatics from "hoist-non-react-statics"
 
 import FullScreenWrapper from "@streamlit/lib/src/components/shared/FullScreenWrapper"
 
-interface Props {
+export interface Props {
   width: number
   height?: number
 }