From cef0dc141ba31241b375be6c4013c7d82ca7f75d Mon Sep 17 00:00:00 2001
From: DimaAmega <dima.amega1@gmail.com>
Date: Fri, 13 Dec 2024 16:04:24 +0400
Subject: [PATCH] UI: refactor use widget wrapper

---
 .../src/components/DashboardContent.tsx       |   8 +-
 .../src/components/DashboardWidgets.tsx       |   6 +-
 .../src/contexts/WidgetWrapper.tsx            |  13 ++
 .../evidently-ui-lib/src/widgets/Widget.tsx   | 199 +++++++++---------
 .../src/widgets/WidgetRenderer.tsx            |   8 +-
 5 files changed, 117 insertions(+), 117 deletions(-)
 create mode 100644 ui/packages/evidently-ui-lib/src/contexts/WidgetWrapper.tsx

diff --git a/ui/packages/evidently-ui-lib/src/components/DashboardContent.tsx b/ui/packages/evidently-ui-lib/src/components/DashboardContent.tsx
index d0c2d0136d..3427c2881e 100644
--- a/ui/packages/evidently-ui-lib/src/components/DashboardContent.tsx
+++ b/ui/packages/evidently-ui-lib/src/components/DashboardContent.tsx
@@ -4,17 +4,13 @@ import { WidgetRenderer } from '~/widgets/WidgetRenderer'
 
 export interface DashboardContentProps {
   widgets: WidgetInfo[]
-  ItemWrapper?: ({ id, children }: { id: string; children: React.ReactNode }) => React.ReactNode
 }
 
-export const DashboardContentWidgets: FunctionComponent<DashboardContentProps> = ({
-  widgets,
-  ItemWrapper
-}) => (
+export const DashboardContentWidgets: FunctionComponent<DashboardContentProps> = ({ widgets }) => (
   <>
     {widgets.length > 0 &&
       widgets.map((wi, idx) => (
-        <React.Fragment key={wi.id}>{WidgetRenderer(`wi_${idx}`, wi, ItemWrapper)}</React.Fragment>
+        <React.Fragment key={wi.id}>{WidgetRenderer(`wi_${idx}`, wi)}</React.Fragment>
       ))}
   </>
 )
diff --git a/ui/packages/evidently-ui-lib/src/components/DashboardWidgets.tsx b/ui/packages/evidently-ui-lib/src/components/DashboardWidgets.tsx
index 05d3e02a11..913943de21 100644
--- a/ui/packages/evidently-ui-lib/src/components/DashboardWidgets.tsx
+++ b/ui/packages/evidently-ui-lib/src/components/DashboardWidgets.tsx
@@ -3,11 +3,9 @@ import type { WidgetInfo } from '~/api'
 import { DashboardContentWidgets } from '~/components/DashboardContent'
 
 export const DashboardWidgets = ({
-  widgets,
-  ItemWrapper
+  widgets
 }: {
   widgets: WidgetInfo[]
-  ItemWrapper?: ({ id, children }: { id: string; children: React.ReactNode }) => React.ReactNode
 }) => {
   if (widgets.length === 0) {
     return (
@@ -20,7 +18,7 @@ export const DashboardWidgets = ({
   return (
     <>
       <Grid container spacing={3} direction='row' alignItems='stretch'>
-        <DashboardContentWidgets widgets={widgets} ItemWrapper={ItemWrapper} />
+        <DashboardContentWidgets widgets={widgets} />
       </Grid>
     </>
   )
diff --git a/ui/packages/evidently-ui-lib/src/contexts/WidgetWrapper.tsx b/ui/packages/evidently-ui-lib/src/contexts/WidgetWrapper.tsx
new file mode 100644
index 0000000000..287ac192d5
--- /dev/null
+++ b/ui/packages/evidently-ui-lib/src/contexts/WidgetWrapper.tsx
@@ -0,0 +1,13 @@
+import React, { useContext } from 'react'
+
+export type WidgetWrapper = {
+  WidgetWrapper: ({ id, children }: { id: string; children: React.ReactNode }) => JSX.Element
+}
+
+const EmptyWidgetWrapper: WidgetWrapper['WidgetWrapper'] = ({ children }) => <>{children}</>
+
+export const widgetWrapperContext = React.createContext<WidgetWrapper>({
+  WidgetWrapper: EmptyWidgetWrapper
+})
+
+export const useWidgetWrapper = () => useContext(widgetWrapperContext)
diff --git a/ui/packages/evidently-ui-lib/src/widgets/Widget.tsx b/ui/packages/evidently-ui-lib/src/widgets/Widget.tsx
index db09f06c1b..5399e97c76 100644
--- a/ui/packages/evidently-ui-lib/src/widgets/Widget.tsx
+++ b/ui/packages/evidently-ui-lib/src/widgets/Widget.tsx
@@ -3,6 +3,7 @@ import type { WidgetInfo } from '~/api'
 
 import { Card, CardContent, Grid, Typography } from '@mui/material'
 
+import { useWidgetWrapper } from '~/contexts/WidgetWrapper'
 import AlertBlock from './AlertBlock'
 import AlertStatBlock from './AlertStatBlock'
 import InsightBlock from './InsightBlock'
@@ -59,127 +60,123 @@ function Sizes(size: number): sizes {
   }
 }
 
-const Widget = (
-  props: WidgetProps & {
-    ItemWrapper?: ({ id, children }: { id: string; children: React.ReactNode }) => React.ReactNode
-  }
-) => {
-  const { size, ItemWrapper } = props
+const Widget = (props: WidgetProps) => {
+  const { size } = props
   const alertsPosition = props.children.alertsPosition ?? 'row'
   const { id, title, details, content, alerts, alertStats, insights } = props.children
   const isAlertsExists = alerts === undefined ? false : alerts.length > 0
   const isInsightsExists = insights === undefined ? false : insights.length > 0
 
-  const Component = (
-    <Card elevation={0}>
-      <CardContent>
-        <Grid container spacing={1} direction={'column'}>
-          {alertsPosition === 'row' ? (
-            <Grid container spacing={1} item>
-              <Grid item xs={isAlertsExists && alertsPosition === 'row' ? 9 : 12}>
-                {title ? (
-                  <Typography fontWeight={500} variant={'h5'}>
-                    {title}
-                  </Typography>
-                ) : (
-                  <div />
-                )}
-                <div>{content}</div>
-                {details ? <Typography variant={'subtitle1'}>{details}</Typography> : <div />}
-              </Grid>
-              {isAlertsExists ? (
-                <Grid
-                  container
-                  spacing={1}
-                  direction={'column'}
-                  justifyContent={'center'}
-                  item
-                  xs={3}
-                >
-                  {alerts ? (
-                    <React.Fragment>
-                      {alertStats ? (
-                        <Grid item>
-                          <AlertStatBlock
-                            alertStats={alertStats}
-                            // classes={classes}
-                          />
-                        </Grid>
+  const { WidgetWrapper } = useWidgetWrapper()
+
+  return (
+    <Grid item {...Sizes(size)}>
+      <WidgetWrapper id={id}>
+        <Card elevation={0}>
+          <CardContent>
+            <Grid container spacing={1} direction={'column'}>
+              {alertsPosition === 'row' ? (
+                <Grid container spacing={1} item>
+                  <Grid item xs={isAlertsExists && alertsPosition === 'row' ? 9 : 12}>
+                    {title ? (
+                      <Typography fontWeight={500} variant={'h5'}>
+                        {title}
+                      </Typography>
+                    ) : (
+                      <div />
+                    )}
+                    <div>{content}</div>
+                    {details ? <Typography variant={'subtitle1'}>{details}</Typography> : <div />}
+                  </Grid>
+                  {isAlertsExists ? (
+                    <Grid
+                      container
+                      spacing={1}
+                      direction={'column'}
+                      justifyContent={'center'}
+                      item
+                      xs={3}
+                    >
+                      {alerts ? (
+                        <React.Fragment>
+                          {alertStats ? (
+                            <Grid item>
+                              <AlertStatBlock
+                                alertStats={alertStats}
+                                // classes={classes}
+                              />
+                            </Grid>
+                          ) : (
+                            <div />
+                          )}
+                          {alerts.map((alert) => (
+                            // biome-ignore lint/correctness/useJsxKeyInIterable: not reordered
+                            <Grid item>
+                              <AlertBlock data={alert} />
+                            </Grid>
+                          ))}
+                        </React.Fragment>
                       ) : (
                         <div />
                       )}
-                      {alerts.map((alert) => (
-                        // biome-ignore lint/correctness/useJsxKeyInIterable: not reordered
-                        <Grid item>
-                          <AlertBlock data={alert} />
-                        </Grid>
-                      ))}
-                    </React.Fragment>
+                    </Grid>
                   ) : (
                     <div />
                   )}
                 </Grid>
               ) : (
-                <div />
-              )}
-            </Grid>
-          ) : (
-            <React.Fragment>
-              <Grid item>
-                {title ? <Typography variant={'h5'}>{title}</Typography> : <div />}
-                <div>{content}</div>
-                {details ? <Typography variant={'subtitle1'}>{details}</Typography> : <div />}
-              </Grid>
-              {isAlertsExists ? (
-                <Grid item xs>
-                  <Grid container direction={'row'} spacing={1}>
-                    {alerts ? (
-                      <React.Fragment>
-                        {alertStats ? (
-                          <Grid item xs>
-                            <AlertStatBlock
-                              alertStats={alertStats}
-                              // classes={classes}
-                            />
-                          </Grid>
+                <React.Fragment>
+                  <Grid item>
+                    {title ? <Typography variant={'h5'}>{title}</Typography> : <div />}
+                    <div>{content}</div>
+                    {details ? <Typography variant={'subtitle1'}>{details}</Typography> : <div />}
+                  </Grid>
+                  {isAlertsExists ? (
+                    <Grid item xs>
+                      <Grid container direction={'row'} spacing={1}>
+                        {alerts ? (
+                          <React.Fragment>
+                            {alertStats ? (
+                              <Grid item xs>
+                                <AlertStatBlock
+                                  alertStats={alertStats}
+                                  // classes={classes}
+                                />
+                              </Grid>
+                            ) : (
+                              <div />
+                            )}
+                            {alerts.map((alert) => (
+                              // biome-ignore lint/correctness/useJsxKeyInIterable: not reordered
+                              <Grid item xs>
+                                <AlertBlock data={alert} />
+                              </Grid>
+                            ))}
+                          </React.Fragment>
                         ) : (
                           <div />
                         )}
-                        {alerts.map((alert) => (
-                          // biome-ignore lint/correctness/useJsxKeyInIterable: not reordered
-                          <Grid item xs>
-                            <AlertBlock data={alert} />
-                          </Grid>
-                        ))}
-                      </React.Fragment>
-                    ) : (
-                      <div />
-                    )}
+                      </Grid>
+                    </Grid>
+                  ) : (
+                    <div />
+                  )}
+                </React.Fragment>
+              )}
+              {isInsightsExists ? (
+                insights?.map((insight) => (
+                  // biome-ignore lint/correctness/useJsxKeyInIterable: not reordered
+                  <Grid item xs sm md>
+                    <InsightBlock data={insight} />
                   </Grid>
-                </Grid>
+                ))
               ) : (
                 <div />
               )}
-            </React.Fragment>
-          )}
-          {isInsightsExists ? (
-            insights?.map((insight) => (
-              // biome-ignore lint/correctness/useJsxKeyInIterable: not reordered
-              <Grid item xs sm md>
-                <InsightBlock data={insight} />
-              </Grid>
-            ))
-          ) : (
-            <div />
-          )}
-        </Grid>
-      </CardContent>
-    </Card>
-  )
-
-  return (
-    <Grid item {...Sizes(size)}>
-      {ItemWrapper ? ItemWrapper({ id, children: Component }) : Component}
+            </Grid>
+          </CardContent>
+        </Card>
+      </WidgetWrapper>
     </Grid>
   )
 }
diff --git a/ui/packages/evidently-ui-lib/src/widgets/WidgetRenderer.tsx b/ui/packages/evidently-ui-lib/src/widgets/WidgetRenderer.tsx
index c0713206cb..587ccc99d4 100644
--- a/ui/packages/evidently-ui-lib/src/widgets/WidgetRenderer.tsx
+++ b/ui/packages/evidently-ui-lib/src/widgets/WidgetRenderer.tsx
@@ -46,11 +46,7 @@ function sizeTransform(size: WidgetSize): 1 | 3 | 6 | 12 {
   return 12
 }
 
-export function WidgetRenderer(
-  key: string,
-  info: WidgetInfo,
-  ItemWrapper?: ({ id, children }: { id: string; children: React.ReactNode }) => React.ReactNode
-) {
+export function WidgetRenderer(key: string, info: WidgetInfo) {
   let content = <NotImplementedWidgetContent />
   if (info.type === 'counter') {
     content = <CounterWidgetContent {...(info.params as CounterWidgetParams)} />
@@ -102,7 +98,7 @@ export function WidgetRenderer(
     content = <TestSuiteWidgetContent {...(info.params as TestSuiteWidgetParams)} />
   }
   return (
-    <Widget key={key} size={sizeTransform(info.size)} ItemWrapper={ItemWrapper}>
+    <Widget key={key} size={sizeTransform(info.size)}>
       {{
         ...info,
         content: content