export function withStreamlitConnection()

in src/streamlit/StreamlitReact.tsx [70:153]
58 lines of code
3 McCabe index (conditional complexity)

export function withStreamlitConnection(
  WrappedComponent: React.ComponentType<ComponentProps>
): React.ComponentType {
  interface WrapperProps {}

  interface WrapperState {
    renderData?: RenderData
    componentError?: Error
  }

  class ComponentWrapper extends React.PureComponent<
    WrapperProps,
    WrapperState
  > {
    public constructor(props: WrapperProps) {
      super(props)
      this.state = {
        renderData: undefined,
        componentError: undefined,
      }
    }

    public static getDerivedStateFromError = (
      error: Error
    ): Partial<WrapperState> => {
      return { componentError: error }
    }

    public componentDidMount = (): void => {
      // Set up event listeners, and signal to Streamlit that we're ready.
      // We won't render the component until we receive the first RENDER_EVENT.
      Streamlit.events.addEventListener(
        Streamlit.RENDER_EVENT,
        this.onRenderEvent
      )
      Streamlit.setComponentReady()
    }

    public componentWillUnmount = (): void => {
      Streamlit.events.removeEventListener(
        Streamlit.RENDER_EVENT,
        this.onRenderEvent
      )
    }

    /**
     * Streamlit is telling this component to redraw.
     * We save the render data in State, so that it can be passed to the
     * component in our own render() function.
     */
    private onRenderEvent = (event: Event): void => {
      // Update our state with the newest render data
      const renderEvent = event as CustomEvent<RenderData>
      this.setState({ renderData: renderEvent.detail })
    }

    public render = (): ReactNode => {
      // If our wrapped component threw an error, display it.
      if (this.state.componentError != null) {
        return (
          <div>
            <h1>Component Error</h1>
            <span>{this.state.componentError.message}</span>
          </div>
        )
      }

      // Don't render until we've gotten our first RENDER_EVENT from Streamlit.
      if (this.state.renderData == null) {
        return null
      }

      return (
        <WrappedComponent
          width={window.innerWidth}
          disabled={this.state.renderData.disabled}
          args={this.state.renderData.args}
        />
      )
    }
  }

  return hoistNonReactStatics(ComponentWrapper, WrappedComponent)
}