feat: Add first class Image support for PIL images (#2113)

* init * improved test and initial implementation * initial work done * python work * UI side in prog * little more python work * some more ui work * make test images bigger * fix images * fix drawer resize * broken object viewer * broken object viewer 2 * got object viewer working * ok, got datasets done i believ * more tests * inital lint * init * a bunch more entity handling * remove sidebar stuff from this PR * little comments * A bunch of little changes * comment * first improvement * little style * lint * finished base implementation of the UI * a little styling * Comment - Checkpoint for Context * Checkpoint part 1: more code * Checkpoint part 1: more code * Checkpoint part 2: Change to context * Comments * Soften table support * removes load calls from UI * possibly done * Done for real besides code cleaning * Done for real besides code cleaning 2 * Lint * Lint * Comment clean * Comment clean 2 * added comment
wandb · Aug 15, 2024 · 3a270ac · 3a270ac
1 parent 930abbd
commit 3a270ac
Show file tree

Hide file tree

Showing 26 changed files with 631 additions and 73 deletions.
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse2/CellValue.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse2/CellValue.tsx
@@ -6,6 +6,8 @@ import {parseRef} from '../../../../react';
 import {ValueViewNumber} from '../Browse3/pages/CallPage/ValueViewNumber';
 import {ValueViewPrimitive} from '../Browse3/pages/CallPage/ValueViewPrimitive';
 import {isRef} from '../Browse3/pages/common/util';
+import {isCustomWeaveTypePayload} from '../Browse3/typeViews/customWeaveType.types';
+import {CustomWeaveTypeDispatcher} from '../Browse3/typeViews/CustomWeaveTypeDispatcher';
 import {CellValueBoolean} from './CellValueBoolean';
 import {CellValueImage} from './CellValueImage';
 import {CellValueString} from './CellValueString';
@@ -64,5 +66,8 @@ export const CellValue = ({value, isExpanded = false}: CellValueProps) => {
       </Box>
     );
   }
+  if (isCustomWeaveTypePayload(value)) {
+    return <CustomWeaveTypeDispatcher data={value} />;
+  }
   return <CellValueString value={JSON.stringify(value)} />;
 };
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse2/browse2Util.ts b/weave-js/src/components/PagePanelComponents/Home/Browse2/browse2Util.ts
@@ -1,9 +1,27 @@
-export const flattenObject = (
+/**
+ * Flatten an object, but preserve any object that has a `_type` field.
+ * This is critical for handling "Weave Types" - payloads that should be
+ * treated as holistic objects, rather than flattened.
+ */
+export const flattenObjectPreservingWeaveTypes = (obj: {
+  [key: string]: any;
+}) => {
+  return flattenObject(obj, '', {}, (key, value) => {
+    return typeof value !== 'object' || value == null || value._type == null;
+  });
+};
+
+const flattenObject = (
   obj: {[key: string]: any},
   parentKey: string = '',
-  result: {[key: string]: any} = {}
+  result: {[key: string]: any} = {},
+  shouldFlatten: (key: string, value: any) => boolean = () => true
 ) => {
-  if (typeof obj !== 'object' || obj === null) {
+  if (
+    typeof obj !== 'object' ||
+    obj === null ||
+    !shouldFlatten(parentKey, obj)
+  ) {
     return obj;
   }
   const keys = Object.keys(obj);
@@ -14,31 +32,14 @@ export const flattenObject = (
     const newKey = parentKey ? `${parentKey}.${key}` : key;
     if (Array.isArray(obj[key])) {
       result[newKey] = obj[key];
-    } else if (typeof obj[key] === 'object') {
-      flattenObject(obj[key], newKey, result);
+    } else if (
+      typeof obj[key] === 'object' &&
+      shouldFlatten(newKey, obj[key])
+    ) {
+      flattenObject(obj[key], newKey, result, shouldFlatten);
     } else {
       result[newKey] = obj[key];
     }
   });
   return result;
 };
-export const unflattenObject = (obj: {[key: string]: any}) => {
-  const result: {[key: string]: any} = {};
-  for (const key in obj) {
-    if (!obj.hasOwnProperty(key)) {
-      continue;
-    }
-    const keys = key.split('.');
-    let current = result;
-    for (let i = 0; i < keys.length; i++) {
-      const k = keys[i];
-      if (i === keys.length - 1) {
-        current[k] = obj[key];
-      } else {
-        current[k] = current[k] || {};
-      }
-      current = current[k];
-    }
-  }
-  return result;
-};
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/filters/CellFilterWrapper.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/filters/CellFilterWrapper.tsx
@@ -11,6 +11,7 @@ type CellFilterWrapperProps = {
   field: string;
   operation: string | null;
   value: any;
+  style?: React.CSSProperties;
 };
 
 export const CellFilterWrapper = ({
@@ -19,6 +20,7 @@ export const CellFilterWrapper = ({
   field,
   operation,
   value,
+  style,
 }: CellFilterWrapperProps) => {
   const onClickCapture = onAddFilter
     ? (e: React.MouseEvent) => {
@@ -31,5 +33,9 @@ export const CellFilterWrapper = ({
       }
     : undefined;
 
-  return <div onClickCapture={onClickCapture}>{children}</div>;
+  return (
+    <div style={style ?? {}} onClickCapture={onClickCapture}>
+      {children}
+    </div>
+  );
 };
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/CallDetails.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/CallDetails.tsx
@@ -8,6 +8,7 @@ import styled from 'styled-components';
 import {MOON_800} from '../../../../../../common/css/color.styles';
 import {Button} from '../../../../../Button';
 import {useWeaveflowRouteContext, WeaveflowPeekContext} from '../../context';
+import {CustomWeaveTypeProjectContext} from '../../typeViews/CustomWeaveTypeDispatcher';
 import {CallsTable} from '../CallsPage/CallsTable';
 import {KeyValueTable} from '../common/KeyValueTable';
 import {CallLink, opNiceName} from '../common/Links';
@@ -117,7 +118,10 @@ export const CallDetails: FC<{
             flex: '0 0 auto',
             p: 2,
           }}>
-          <ObjectViewerSection title="Inputs" data={inputs} />
+          <CustomWeaveTypeProjectContext.Provider
+            value={{entity: call.entity, project: call.project}}>
+            <ObjectViewerSection title="Inputs" data={inputs} />
+          </CustomWeaveTypeProjectContext.Provider>
         </Box>
         <Box
           sx={{
@@ -132,7 +136,10 @@ export const CallDetails: FC<{
               <ExceptionDetails exceptionInfo={excInfo} />
             </>
           ) : (
-            <ObjectViewerSection title="Outputs" data={output} />
+            <CustomWeaveTypeProjectContext.Provider
+              value={{entity: call.entity, project: call.project}}>
+              <ObjectViewerSection title="Output" data={output} />
+            </CustomWeaveTypeProjectContext.Provider>
           )}
         </Box>
         {multipleChildCallOpRefs.map(opVersionRef => {
@@ -251,13 +258,15 @@ const getDisplayInputsAndOutput = (call: CallSchema) => {
   const span = call.rawSpan;
   const inputKeys =
     span.inputs._keys ??
-    Object.keys(span.inputs).filter(k => !k.startsWith('_'));
+    Object.keys(span.inputs).filter(k => !k.startsWith('_') || k === '_type');
   const inputs = _.fromPairs(inputKeys.map(k => [k, span.inputs[k]]));
 
   const callOutput = span.output ?? {};
   const outputKeys =
     callOutput._keys ??
-    Object.keys(callOutput).filter(k => k === '_result' || !k.startsWith('_'));
+    Object.keys(callOutput).filter(
+      k => k === '_result' || !k.startsWith('_') || k === '_type'
+    );
   const output = _.fromPairs(outputKeys.map(k => [k, callOutput[k]]));
   return {inputs, output};
 };

diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/DataTableView.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/DataTableView.tsx
@@ -13,14 +13,19 @@ import _ from 'lodash';
 import React, {FC, useCallback, useContext, useEffect, useMemo} from 'react';
 import {useHistory} from 'react-router-dom';
 
-import {isWeaveObjectRef, parseRef} from '../../../../../../react';
-import {flattenObject} from '../../../Browse2/browse2Util';
+import {
+  isWeaveObjectRef,
+  parseRef,
+  WeaveObjectRef,
+} from '../../../../../../react';
+import {flattenObjectPreservingWeaveTypes} from '../../../Browse2/browse2Util';
 import {CellValue} from '../../../Browse2/CellValue';
 import {
   useWeaveflowCurrentRouteContext,
   WeaveflowPeekContext,
 } from '../../context';
 import {StyledDataGrid} from '../../StyledDataGrid';
+import {CustomWeaveTypeProjectContext} from '../../typeViews/CustomWeaveTypeDispatcher';
 import {TABLE_ID_EDGE_NAME} from '../wfReactInterface/constants';
 import {useWFHooks} from '../wfReactInterface/context';
 import {TableQuery} from '../wfReactInterface/wfDataModelHooksInterface';
@@ -53,6 +58,11 @@ export const WeaveCHTable: FC<{
     limit: MAX_ROWS + 1,
   });
 
+  const parsedRef = useMemo(
+    () => parseRef(props.tableRefUri) as WeaveObjectRef,
+    [props.tableRefUri]
+  );
+
   // Determines if the table itself is truncated
   const isTruncated = useMemo(() => {
     return (fetchQuery.result ?? []).length > MAX_ROWS;
@@ -96,16 +106,19 @@ export const WeaveCHTable: FC<{
   );
 
   return (
-    <DataTableView
-      data={sourceRows ?? []}
-      loading={fetchQuery.loading}
-      isTruncated={isTruncated}
-      // Display key is "val" as the resulting rows have metadata/ref
-      // information outside of the actual data
-      displayKey="val"
-      onLinkClick={onClickEnabled ? onClick : undefined}
-      fullHeight={props.fullHeight}
-    />
+    <CustomWeaveTypeProjectContext.Provider
+      value={{entity: parsedRef.entityName, project: parsedRef.projectName}}>
+      <DataTableView
+        data={sourceRows ?? []}
+        loading={fetchQuery.loading}
+        isTruncated={isTruncated}
+        // Display key is "val" as the resulting rows have metadata/ref
+        // information outside of the actual data
+        displayKey="val"
+        onLinkClick={onClickEnabled ? onClick : undefined}
+        fullHeight={props.fullHeight}
+      />
+    </CustomWeaveTypeProjectContext.Provider>
   );
 };
 
@@ -133,7 +146,7 @@ export const DataTableView: FC<{
       if (val == null) {
         return {};
       } else if (typeof val === 'object' && !Array.isArray(val)) {
-        return flattenObject(val);
+        return flattenObjectPreservingWeaveTypes(val);
       }
       return {'': val};
     });

diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/ObjectViewer.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/ObjectViewer.tsx
@@ -21,6 +21,7 @@ import {LoadingDots} from '../../../../../LoadingDots';
 import {Browse2OpDefCode} from '../../../Browse2/Browse2OpDefCode';
 import {parseRefMaybe} from '../../../Browse2/SmallRef';
 import {StyledDataGrid} from '../../StyledDataGrid';
+import {isCustomWeaveTypePayload} from '../../typeViews/customWeaveType.types';
 import {isRef} from '../common/util';
 import {
   LIST_INDEX_EDGE_NAME,
@@ -163,6 +164,37 @@ export const ObjectViewer = ({
       }
     > = [];
     traverse(resolvedData, context => {
+      // Ops should be migrated to the generic CustomWeaveType pattern, but for
+      // now they are custom handled.
+      const isOpPayload = context.value?.weave_type?.type === 'Op';
+
+      if (isCustomWeaveTypePayload(context.value) && !isOpPayload) {
+        /**
+         * This block adds an "empty" key that is used to render the custom
+         * weave type. In the event that a custom type has both properties AND
+         * custom views, then we might need to extend / modify this part.
+         */
+        const refBackingData = context.value?._ref;
+        let depth = context.depth;
+        let path = context.path;
+        if (refBackingData) {
+          contexts.push({
+            ...context,
+            isExpandableRef: true,
+          });
+          depth += 1;
+          path = context.path.plus('');
+        }
+        contexts.push({
+          depth,
+          isLeaf: true,
+          path,
+          value: context.value,
+          valueType: context.valueType,
+        });
+        return 'skip';
+      }
+
       if (context.depth !== 0) {
         const contextTail = context.path.tail();
         const isNullDescription =
@@ -207,7 +239,8 @@ export const ObjectViewer = ({
       if (USE_TABLE_FOR_ARRAYS && context.valueType === 'array') {
         return 'skip';
       }
-      if (context.value?._ref && context.value?.weave_type?.type === 'Op') {
+      if (context.value?._ref && isOpPayload) {
+        // This should be moved to the CustomWeaveType pattern.
         contexts.push({
           depth: context.depth + 1,
           isLeaf: true,
@@ -377,11 +410,15 @@ export const ObjectViewer = ({
             isRef(params.model.value) &&
             (parseRefMaybe(params.model.value) as any).weaveKind === 'table';
           const {isCode} = params.model;
+          const isCustomWeaveType = isCustomWeaveTypePayload(
+            params.model.value
+          );
           if (
             isNonRefString ||
             (isArray && USE_TABLE_FOR_ARRAYS) ||
             isTableRef ||
-            isCode
+            isCode ||
+            isCustomWeaveType
           ) {
             return 'auto';
           }

diff --git a/...js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/ObjectViewerSection.tsx b/...js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/ObjectViewerSection.tsx
@@ -14,6 +14,8 @@ import {isWeaveObjectRef, parseRef} from '../../../../../../react';
 import {Alert} from '../../../../../Alert';
 import {Button} from '../../../../../Button';
 import {CodeEditor} from '../../../../../CodeEditor';
+import {isCustomWeaveTypePayload} from '../../typeViews/customWeaveType.types';
+import {CustomWeaveTypeDispatcher} from '../../typeViews/CustomWeaveTypeDispatcher';
 import {isRef} from '../common/util';
 import {OBJECT_ATTR_EDGE_NAME} from '../wfReactInterface/constants';
 import {WeaveCHTable, WeaveCHTableSourceRefContext} from './DataTableView';
@@ -119,7 +121,7 @@ const ObjectViewerSectionNonEmpty = ({
       );
     }
     return null;
-  }, [apiRef, mode, data, expandedIds]);
+  }, [mode, apiRef, data, expandedIds]);
 
   const setTreeExpanded = useCallback(
     (setIsExpanded: boolean) => {
@@ -215,9 +217,20 @@ export const ObjectViewerSection = ({
   noHide,
   isExpanded,
 }: ObjectViewerSectionProps) => {
-  const numKeys = Object.keys(data).length;
   const currentRef = useContext(WeaveCHTableSourceRefContext);
 
+  if (isCustomWeaveTypePayload(data)) {
+    return (
+      <>
+        <TitleRow>
+          <Title>{title}</Title>
+        </TitleRow>
+        <CustomWeaveTypeDispatcher data={data} />
+      </>
+    );
+  }
+
+  const numKeys = Object.keys(data).length;
   if (numKeys === 0) {
     return (
       <>

diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/ValueView.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/ValueView.tsx
@@ -1,7 +1,9 @@
 import React, {useMemo} from 'react';
 
-import {parseRef} from '../../../../../../react';
+import {isWeaveObjectRef, parseRef} from '../../../../../../react';
 import {parseRefMaybe, SmallRef} from '../../../Browse2/SmallRef';
+import {isCustomWeaveTypePayload} from '../../typeViews/customWeaveType.types';
+import {CustomWeaveTypeDispatcher} from '../../typeViews/CustomWeaveTypeDispatcher';
 import {isRef} from '../common/util';
 import {
   DataTableView,
@@ -77,5 +79,36 @@ export const ValueView = ({data, isExpanded}: ValueViewProps) => {
     return <div>{JSON.stringify(data.value)}</div>;
   }
 
+  if (data.valueType === 'object') {
+    if (isCustomWeaveTypePayload(data.value)) {
+      // This is a little ugly, but essentially if the data is coming from an
+      // expanded ref, then we want to use that ref to get the entity and project.
+      // Else we just use the current entity and project.
+      let entityForWeaveType: string | undefined;
+      let projectForWeaveType: string | undefined;
+
+      if (valueIsExpandedRef(data)) {
+        const parsedRef = parseRef((data.value as any)._ref);
+        if (isWeaveObjectRef(parsedRef)) {
+          entityForWeaveType = parsedRef.entityName;
+          projectForWeaveType = parsedRef.projectName;
+        }
+      }
+
+      // If we have have a custom view for this weave type, use it.
+      return (
+        <CustomWeaveTypeDispatcher
+          data={data.value}
+          entity={entityForWeaveType}
+          project={projectForWeaveType}
+        />
+      );
+    }
+  }
+
   return <div>{data.value.toString()}</div>;
 };
+
+const valueIsExpandedRef = (data: ValueData) => {
+  return data.value != null && (data.value as any)._ref != null;
+};