feat: support pivots menu (#1755)

packages/component/src/ui/tree-view/TreeNode.tsx

@@ -1,4 +1,4 @@
-import { useEffect, useState } from 'react';
+import { useEffect } from 'react';
 import { useDrag, useDrop } from 'react-dnd';
 
 import {
@@ -14,21 +14,21 @@ import type {
   TreeNodeProps,
 } from './types';
 
-const NodeLine = <N,>({
+const NodeLine = <RenderProps,>({
   node,
   onDrop,
   allowDrop = true,
   isTop = false,
-}: NodeLIneProps<N>) => {
+}: NodeLIneProps<RenderProps>) => {
   const [{ isOver }, drop] = useDrop(
     () => ({
       accept: 'node',
-      drop: (item: Node<N>, monitor) => {
+      drop: (item: Node<RenderProps>, monitor) => {
         const didDrop = monitor.didDrop();
         if (didDrop) {
           return;
         }
-        onDrop?.(item, node, {
+        onDrop?.(item.id, node.id, {
           internal: false,
           topLine: isTop,
           bottomLine: !isTop,
@@ -44,24 +44,23 @@ const NodeLine = <N,>({
 
   return <StyledNodeLine ref={drop} show={isOver && allowDrop} isTop={isTop} />;
 };
-const TreeNodeItem = <N,>({
+const TreeNodeItemWithDnd = <RenderProps,>({
   node,
   allowDrop,
-  collapsed,
   setCollapsed,
   ...otherProps
-}: TreeNodeItemProps<N>) => {
+}: TreeNodeItemProps<RenderProps>) => {
   const { onAdd, onDelete, onDrop } = otherProps;
 
   const [{ canDrop, isOver }, drop] = useDrop(
     () => ({
       accept: 'node',
-      drop: (item: Node<N>, monitor) => {
+      drop: (item: Node<RenderProps>, monitor) => {
         const didDrop = monitor.didDrop();
         if (didDrop || item.id === node.id || !allowDrop) {
           return;
         }
-        onDrop?.(item, node, {
+        onDrop?.(item.id, node.id, {
           internal: true,
           topLine: false,
           bottomLine: false,
@@ -77,44 +76,79 @@ const TreeNodeItem = <N,>({
 
   useEffect(() => {
     if (isOver && canDrop) {
-      setCollapsed(false);
+      setCollapsed(node.id, false);
     }
   }, [isOver, canDrop]);
 
   return (
-    <div ref={drop}>
+    <TreeNodeItem
+      dropRef={drop}
+      onAdd={onAdd}
+      onDelete={onDelete}
+      node={node}
+      allowDrop={allowDrop}
+      setCollapsed={setCollapsed}
+      isOver={isOver}
+      canDrop={canDrop}
+      {...otherProps}
+    />
+  );
+};
+
+const TreeNodeItem = <RenderProps,>({
+  node,
+  collapsed,
+  setCollapsed,
+  selectedId,
+  isOver = false,
+  canDrop = false,
+  onAdd,
+  onDelete,
+  dropRef,
+}: TreeNodeItemProps<RenderProps>) => {
+  return (
+    <div ref={dropRef}>
       {node.render?.(node, {
-        isOver: !!(isOver && canDrop),
+        isOver: isOver && canDrop,
         onAdd: () => onAdd?.(node),
         onDelete: () => onDelete?.(node),
         collapsed,
         setCollapsed,
+        isSelected: selectedId === node.id,
       })}
     </div>
   );
 };
 
-export const TreeNode = <N,>({
-  node,
-  index,
-  allowDrop = true,
-  ...otherProps
-}: TreeNodeProps<N>) => {
-  const { indent } = otherProps;
-  const [collapsed, setCollapsed] = useState(false);
-
+export const TreeNodeWithDnd = <RenderProps,>(
+  props: TreeNodeProps<RenderProps>
+) => {
   const [{ isDragging }, drag] = useDrag(() => ({
     type: 'node',
-    item: node,
+    item: props.node,
     collect: monitor => ({
       isDragging: monitor.isDragging(),
     }),
   }));
 
+  return <TreeNode dragRef={drag} isDragging={isDragging} {...props} />;
+};
+
+export const TreeNode = <RenderProps,>({
+  node,
+  index,
+  isDragging = false,
+  allowDrop = true,
+  dragRef,
+  ...otherProps
+}: TreeNodeProps<RenderProps>) => {
+  const { indent, enableDnd, collapsedIds } = otherProps;
+  const collapsed = collapsedIds.includes(node.id);
+
   return (
-    <StyledTreeNodeContainer ref={drag} isDragging={isDragging}>
+    <StyledTreeNodeContainer ref={dragRef} isDragging={isDragging}>
       <StyledTreeNodeWrapper>
-        {index === 0 && (
+        {enableDnd && index === 0 && (
           <NodeLine
             node={node}
             {...otherProps}
@@ -122,15 +156,25 @@ export const TreeNode = <N,>({
             isTop={true}
           />
         )}
-        <TreeNodeItem
-          node={node}
-          index={index}
-          allowDrop={allowDrop}
-          collapsed={collapsed}
-          setCollapsed={setCollapsed}
-          {...otherProps}
-        />
-        {(!node.children?.length || collapsed) && (
+        {enableDnd ? (
+          <TreeNodeItemWithDnd
+            node={node}
+            index={index}
+            allowDrop={allowDrop}
+            collapsed={collapsed}
+            {...otherProps}
+          />
+        ) : (
+          <TreeNodeItem
+            node={node}
+            index={index}
+            allowDrop={allowDrop}
+            collapsed={collapsed}
+            {...otherProps}
+          />
+        )}
+
+        {enableDnd && (!node.children?.length || collapsed) && (
           <NodeLine
             node={node}
             {...otherProps}
@@ -140,18 +184,26 @@ export const TreeNode = <N,>({
       </StyledTreeNodeWrapper>
       <StyledCollapse in={!collapsed} indent={indent}>
         {node.children &&
-          node.children.map((childNode, index) => (
-            <TreeNode
-              key={childNode.id}
-              node={childNode}
-              index={index}
-              allowDrop={isDragging ? false : allowDrop}
-              {...otherProps}
-            />
-          ))}
+          node.children.map((childNode, index) =>
+            enableDnd ? (
+              <TreeNodeWithDnd
+                key={childNode.id}
+                node={childNode}
+                index={index}
+                allowDrop={isDragging ? false : allowDrop}
+                {...otherProps}
+              />
+            ) : (
+              <TreeNode
+                key={childNode.id}
+                node={childNode}
+                index={index}
+                allowDrop={false}
+                {...otherProps}
+              />
+            )
+          )}
       </StyledCollapse>
     </StyledTreeNodeContainer>
   );
 };
-
-export default TreeNode;

packages/component/src/ui/tree-view/TreeView.tsx

@@ -1,15 +1,108 @@
+import { useEffect, useState } from 'react';
 import { DndProvider } from 'react-dnd';
 import { HTML5Backend } from 'react-dnd-html5-backend';
 
-import { TreeNode } from './TreeNode';
-import type { TreeViewProps } from './types';
-export const TreeView = <N,>({ data, ...otherProps }: TreeViewProps<N>) => {
+import { TreeNode, TreeNodeWithDnd } from './TreeNode';
+import type { TreeNodeProps, TreeViewProps } from './types';
+import { flattenIds } from './utils';
+
+export const TreeView = <RenderProps,>({
+  data,
+  enableKeyboardSelection,
+  onSelect,
+  enableDnd = true,
+  initialCollapsedIds = [],
+  ...otherProps
+}: TreeViewProps<RenderProps>) => {
+  const [selectedId, setSelectedId] = useState<string>();
+  // TODO: should record collapsedIds in localStorage
+  const [collapsedIds, setCollapsedIds] =
+    useState<string[]>(initialCollapsedIds);
+
+  useEffect(() => {
+    if (!enableKeyboardSelection) {
+      return;
+    }
+
+    const flattenedIds = flattenIds<RenderProps>(data);
+
+    const handleDirectionKeyDown = (e: KeyboardEvent) => {
+      if (e.key !== 'ArrowDown' && e.key !== 'ArrowUp') {
+        return;
+      }
+      if (selectedId === undefined) {
+        setSelectedId(flattenedIds[0]);
+        return;
+      }
+      let selectedIndex = flattenedIds.indexOf(selectedId);
+      if (e.key === 'ArrowDown') {
+        selectedIndex < flattenedIds.length - 1 && selectedIndex++;
+      }
+      if (e.key === 'ArrowUp') {
+        selectedIndex > 0 && selectedIndex--;
+      }
+
+      setSelectedId(flattenedIds[selectedIndex]);
+    };
+
+    const handleEnterKeyDown = (e: KeyboardEvent) => {
+      if (e.key !== 'Enter') {
+        return;
+      }
+      selectedId && onSelect?.(selectedId);
+    };
+
+    document.addEventListener('keydown', handleDirectionKeyDown);
+    document.addEventListener('keydown', handleEnterKeyDown);
+
+    return () => {
+      document.removeEventListener('keydown', handleDirectionKeyDown);
+      document.removeEventListener('keydown', handleEnterKeyDown);
+    };
+  }, [data, selectedId]);
+
+  const setCollapsed: TreeNodeProps['setCollapsed'] = (id, collapsed) => {
+    if (collapsed) {
+      setCollapsedIds(ids => [...ids, id]);
+    } else {
+      setCollapsedIds(ids => ids.filter(i => i !== id));
+    }
+  };
+
+  if (enableDnd) {
+    return (
+      <DndProvider backend={HTML5Backend}>
+        {data.map((node, index) => (
+          <TreeNodeWithDnd
+            key={node.id}
+            index={index}
+            collapsedIds={collapsedIds}
+            setCollapsed={setCollapsed}
+            node={node}
+            selectedId={selectedId}
+            enableDnd={enableDnd}
+            {...otherProps}
+          />
+        ))}
+      </DndProvider>
+    );
+  }
+
   return (
-    <DndProvider backend={HTML5Backend}>
+    <>
       {data.map((node, index) => (
-        <TreeNode key={node.id} index={index} node={node} {...otherProps} />
+        <TreeNode
+          key={node.id}
+          index={index}
+          collapsedIds={collapsedIds}
+          setCollapsed={setCollapsed}
+          node={node}
+          selectedId={selectedId}
+          enableDnd={enableDnd}
+          {...otherProps}
+        />
       ))}
-    </DndProvider>
+    </>
   );
 };
 

packages/component/src/ui/tree-view/styles.ts

@@ -15,8 +15,8 @@ export const StyledTreeNodeWrapper = styled('div')(() => {
     position: 'relative',
   };
 });
-export const StyledTreeNodeContainer = styled('div')<{ isDragging: boolean }>(
-  ({ isDragging, theme }) => {
+export const StyledTreeNodeContainer = styled('div')<{ isDragging?: boolean }>(
+  ({ isDragging = false, theme }) => {
     return {
       background: isDragging ? theme.colors.hoverBackground : '',
       // opacity: isDragging ? 0.4 : 1,

packages/component/src/ui/tree-view/types.ts

@@ -1,52 +1,71 @@
-import type { CSSProperties, ReactNode } from 'react';
+import type { CSSProperties, ReactNode, Ref } from 'react';
 
-export type Node<N> = {
+export type DropPosition = {
+  topLine: boolean;
+  bottomLine: boolean;
+  internal: boolean;
+};
+export type OnDrop = (
+  dragId: string,
+  dropId: string,
+  position: DropPosition
+) => void;
+
+export type Node<RenderProps = unknown> = {
   id: string;
-  children?: Node<N>[];
-  render?: (
-    node: Node<N>,
+  children?: Node<RenderProps>[];
+  render: (
+    node: Node<RenderProps>,
     eventsAndStatus: {
       isOver: boolean;
       onAdd: () => void;
       onDelete: () => void;
       collapsed: boolean;
-      setCollapsed: (collapsed: boolean) => void;
+      setCollapsed: (id: string, collapsed: boolean) => void;
+      isSelected: boolean;
     },
-    extendProps?: unknown
+    renderProps?: RenderProps
   ) => ReactNode;
-} & N;
-
-type CommonProps<N> = {
-  indent?: CSSProperties['paddingLeft'];
-  onAdd?: (node: Node<N>) => void;
-  onDelete?: (node: Node<N>) => void;
-  onDrop?: (
-    dragNode: Node<N>,
-    dropNode: Node<N>,
-    position: {
-      topLine: boolean;
-      bottomLine: boolean;
-      internal: boolean;
-    }
-  ) => void;
 };
 
-export type TreeNodeProps<N> = {
-  node: Node<N>;
+type CommonProps<RenderProps = unknown> = {
+  enableDnd?: boolean;
+  enableKeyboardSelection?: boolean;
+  indent?: CSSProperties['paddingLeft'];
+  onAdd?: (node: Node<RenderProps>) => void;
+  onDelete?: (node: Node<RenderProps>) => void;
+  onDrop?: OnDrop;
+  // Only trigger when the enableKeyboardSelection is true
+  onSelect?: (id: string) => void;
+};
+
+export type TreeNodeProps<RenderProps = unknown> = {
+  node: Node<RenderProps>;
   index: number;
+  collapsedIds: string[];
+  setCollapsed: (id: string, collapsed: boolean) => void;
   allowDrop?: boolean;
-} & CommonProps<N>;
+  selectedId?: string;
+  isDragging?: boolean;
+  dragRef?: Ref<HTMLDivElement>;
+} & CommonProps<RenderProps>;
 
-export type TreeNodeItemProps<N> = {
+export type TreeNodeItemProps<RenderProps = unknown> = {
   collapsed: boolean;
-  setCollapsed: (collapsed: boolean) => void;
-} & TreeNodeProps<N>;
+  setCollapsed: (id: string, collapsed: boolean) => void;
 
-export type TreeViewProps<N> = {
-  data: Node<N>[];
-} & CommonProps<N>;
+  isOver?: boolean;
+  canDrop?: boolean;
 
-export type NodeLIneProps<N> = {
+  dropRef?: Ref<HTMLDivElement>;
+} & TreeNodeProps<RenderProps>;
+
+export type TreeViewProps<RenderProps = unknown> = {
+  data: Node<RenderProps>[];
+  initialCollapsedIds?: string[];
+} & CommonProps<RenderProps>;
+
+export type NodeLIneProps<RenderProps = unknown> = {
   allowDrop: boolean;
   isTop?: boolean;
-} & Pick<TreeNodeProps<N>, 'node' | 'onDrop'>;
+} & Pick<TreeNodeProps<RenderProps>, 'node' | 'onDrop'>;

packages/component/src/ui/tree-view/utils.ts

@@ -0,0 +1,18 @@
+import type { Node } from '@affine/component';
+
+export function flattenIds<RenderProps>(arr: Node<RenderProps>[]): string[] {
+  const result: string[] = [];
+
+  function flatten(arr: Node<RenderProps>[]) {
+    for (let i = 0, len = arr.length; i < len; i++) {
+      const item = arr[i];
+      result.push(item.id);
+      if (Array.isArray(item.children)) {
+        flatten(item.children);
+      }
+    }
+  }
+
+  flatten(arr);
+  return result;
+}