Use withHOC for wrapping components in providers

apps/www/src/registry/default/plate-ui/image-element.tsx

@@ -1,5 +1,5 @@
 import React from 'react';
-import { PlateElement, PlateElementProps, Value } from '@udecode/plate-common';
+import { PlateElement, PlateElementProps, Value, withHOC } from '@udecode/plate-common';
 import {
   ELEMENT_IMAGE,
   Image,
@@ -18,12 +18,12 @@ import {
   ResizeHandle,
 } from './resizable';
 
-function InnerImageElement({
+const ImageElement = withHOC(ResizableProvider, ({
   className,
   children,
   nodeProps,
   ...props
-}: PlateElementProps<Value, TImageElement>) {
+}: PlateElementProps<Value, TImageElement>) => {
   const { readOnly, focused, selected, align = 'center' } = useMediaState();
   const width = useResizableStore().get.width();
 
@@ -69,12 +69,7 @@ function InnerImageElement({
       </PlateElement>
     </MediaPopover>
   );
-}
+});
+ImageElement.displayName = 'ImageElement';
 
-export function ImageElement(props: PlateElementProps<Value, TImageElement>) {
-  return (
-    <ResizableProvider>
-      <InnerImageElement {...props} />
-    </ResizableProvider>
-  );
-}
+export { ImageElement };

apps/www/src/registry/default/plate-ui/media-embed-element.tsx

@@ -1,5 +1,5 @@
 import React from 'react';
-import { PlateElement, PlateElementProps, Value } from '@udecode/plate-common';
+import { PlateElement, PlateElementProps, Value, withHOC } from '@udecode/plate-common';
 import {
   ELEMENT_MEDIA_EMBED,
   parseTwitterUrl,
@@ -21,7 +21,7 @@ import {
   ResizeHandle,
 } from './resizable';
 
-const InnerMediaEmbedElement = React.forwardRef<
+const MediaEmbedElement = withHOC(ResizableProvider, React.forwardRef<
   React.ElementRef<typeof PlateElement>,
   PlateElementProps<Value, TMediaEmbedElement>
 >(({ className, children, ...props }, ref) => {
@@ -135,16 +135,7 @@ const InnerMediaEmbedElement = React.forwardRef<
       </PlateElement>
     </MediaPopover>
   );
-});
-InnerMediaEmbedElement.displayName = 'InnerMediaEmbedElement';
-
-const MediaEmbedElement: typeof InnerMediaEmbedElement = React.forwardRef(
-  (props, ref) => (
-    <ResizableProvider>
-      <InnerMediaEmbedElement {...props} ref={ref} />
-    </ResizableProvider>
-  )
-);
+}));
 MediaEmbedElement.displayName = 'MediaEmbedElement';
 
 export { MediaEmbedElement };

packages/core/src/utils/__snapshots__/withHOC.spec.tsx.snap

@@ -0,0 +1,53 @@
+// Jest Snapshot v1, https://goo.gl/fbAQLP
+
+exports[`withHOC renders component with HOC 1`] = `
+<div>
+  <div>
+    {}
+    null
+    <div>
+      {"myProp":"component-prop"}
+      null
+    </div>
+  </div>
+</div>
+`;
+
+exports[`withHOC renders component with HOC and HOC props 1`] = `
+<div>
+  <div>
+    {"myProp":"hoc-prop"}
+    null
+    <div>
+      {"myProp":"component-prop"}
+      null
+    </div>
+  </div>
+</div>
+`;
+
+exports[`withHOC renders component with HOC and HOC ref 1`] = `
+<div>
+  <div>
+    {}
+    {"current":{"myRef":"hoc-ref"}}
+    <div>
+      {"myProp":"component-prop"}
+      null
+    </div>
+  </div>
+</div>
+`;
+
+exports[`withHOC renders component with HOC and component ref 1`] = `
+<div>
+  <div>
+    {}
+    null
+    <div>
+      {"myProp":"component-prop"}
+      {"current":{"myRef":"component-ref"}}
+    </div>
+  </div>
+</div>
+`;

packages/core/src/utils/withHOC.spec.tsx

@@ -0,0 +1,67 @@
+import React, { ReactNode } from 'react';
+import { render } from '@testing-library/react';
+
+import { withHOC } from './withHOC';
+
+/* eslint-disable react/display-name */
+describe('withHOC', () => {
+  type DummyRef = { myRef: string };
+  type DummyProps = { myProp: string };
+
+  const HOC = React.forwardRef<DummyRef, DummyProps & { children: ReactNode }>(
+    ({ children, ...props }, ref) => (
+      <div>
+        {JSON.stringify(props)}
+        {JSON.stringify(ref)}
+        {children}
+      </div>
+    )
+  );
+
+  const Component = React.forwardRef<DummyRef, DummyProps>((props, ref) => (
+    <div>
+      {JSON.stringify(props)}
+      {JSON.stringify(ref)}
+    </div>
+  ));
+
+  it('renders component with HOC', () => {
+    const WithHOC = withHOC(HOC, Component);
+
+    const { container } = render(<WithHOC myProp="component-prop" />);
+
+    expect(container).toHaveTextContent('component-prop');
+    expect(container).toMatchSnapshot();
+  });
+
+  it('renders component with HOC and HOC props', () => {
+    const WithHOC = withHOC(HOC, Component, { myProp: 'hoc-prop' });
+
+    const { container } = render(<WithHOC myProp="component-prop" />);
+
+    expect(container).toHaveTextContent('hoc-prop');
+    expect(container).toMatchSnapshot();
+  });
+
+  it('renders component with HOC and HOC ref', () => {
+    const hocRef = { current: { myRef: 'hoc-ref' } };
+    const WithHOC = withHOC(HOC, Component, undefined, hocRef);
+
+    const { container } = render(<WithHOC myProp="component-prop" />);
+
+    expect(container).toHaveTextContent('hoc-ref');
+    expect(container).toMatchSnapshot();
+  });
+
+  it('renders component with HOC and component ref', () => {
+    const componentRef = { current: { myRef: 'component-ref' } };
+    const WithHOC = withHOC(HOC, Component);
+
+    const { container } = render(
+      <WithHOC myProp="component-prop" ref={componentRef} />
+    );
+
+    expect(container).toHaveTextContent('component-ref');
+    expect(container).toMatchSnapshot();
+  });
+});

packages/core/src/utils/withHOC.tsx

@@ -1,14 +1,18 @@
-import React, { FunctionComponent } from 'react';
+import React from 'react';
 
-export const withHOC = <T,>(
-  HOC: FunctionComponent<any>,
-  Component: FunctionComponent<T>,
-  hocProps?: any
-): FunctionComponent<T> =>
-  function hoc(props: T) {
-    return (
-      <HOC {...hocProps}>
-        <Component {...(props as any)} />
-      </HOC>
-    );
-  };
+type RefComponent<P, R> = React.FC<P> & {
+  ref?: React.Ref<R>;
+};
+
+/* eslint-disable react/display-name */
+export const withHOC = <ComponentProps, HOCProps, ComponentRef, HOCRef>(
+  HOC: RefComponent<HOCProps, HOCRef>,
+  Component: RefComponent<ComponentProps, ComponentRef>,
+  hocProps?: Omit<HOCProps, 'children'>,
+  hocRef?: React.Ref<HOCRef>
+) =>
+  React.forwardRef<ComponentRef, ComponentProps>((props, componentRef) => (
+    <HOC {...(hocProps as any)} ref={hocRef}>
+      <Component {...props} ref={componentRef} />
+    </HOC>
+  ));