fix memoWithAs types

reach · Jun 24, 2020 · 82779c2 · 82779c2
1 parent 22b51e8
commit 82779c2
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 6 deletions.
diff --git a/packages/utils/src/index.tsx b/packages/utils/src/index.tsx
@@ -17,7 +17,10 @@ import {
   ComponentWithAs,
   ComponentWithForwardedRef,
   DistributiveOmit,
+  ElementByTag,
+  ElementTagNameMap,
   ForwardRefExoticComponentWithAs,
+  ForwardRefWithAsRenderFunction,
   FunctionComponentWithAs,
   MemoExoticComponentWithAs,
   PropsFromAs,
@@ -203,12 +206,9 @@ export function createNamedContext<ContextValueType>(
  * the time time being.
  */
 export function forwardRefWithAs<Props, ComponentType extends As = "div">(
-  render: (
-    props: PropsFromAs<ComponentType, Props>,
-    ref: React.RefObject<any>
-  ) => React.ReactElement | null
+  render: ForwardRefWithAsRenderFunction<ComponentType, Props>
 ) {
-  return React.forwardRef(render as any) as ForwardRefExoticComponentWithAs<
+  return React.forwardRef(render) as ForwardRefExoticComponentWithAs<
     ComponentType,
     Props
   >;
@@ -708,6 +708,8 @@ export {
   ComponentWithAs,
   ComponentWithForwardedRef,
   DistributiveOmit,
+  ElementByTag,
+  ElementTagNameMap,
   ForwardRefExoticComponentWithAs,
   FunctionComponentWithAs,
   MemoExoticComponentWithAs,

diff --git a/packages/utils/src/types.tsx b/packages/utils/src/types.tsx
@@ -176,9 +176,56 @@ export interface MemoExoticComponentWithAs<
   ComponentType extends As,
   ComponentProps
 > extends NamedExoticComponentWithAs<ComponentType, ComponentProps> {
-  readonly type: ComponentType;
+  readonly type: ComponentType extends React.ComponentType
+    ? ComponentType
+    : FunctionComponentWithAs<ComponentType, ComponentProps>;
 }
 
+export interface ForwardRefWithAsRenderFunction<
+  ComponentType extends As,
+  ComponentProps = {}
+> {
+  (
+    props: React.PropsWithChildren<PropsFromAs<ComponentType, ComponentProps>>,
+    ref:
+      | ((
+          instance:
+            | (ComponentType extends keyof ElementTagNameMap
+                ? ElementByTag<ComponentType>
+                : any)
+            | null
+        ) => void)
+      | React.MutableRefObject<
+          | (ComponentType extends keyof ElementTagNameMap
+              ? ElementByTag<ComponentType>
+              : any)
+          | null
+        >
+      | null
+  ): React.ReactElement | null;
+  displayName?: string;
+  // explicit rejected with `never` required due to
+  // https://github.com/microsoft/TypeScript/issues/36826
+  /**
+   * defaultProps are not supported on render functions
+   */
+  defaultProps?: never;
+  /**
+   * propTypes are not supported on render functions
+   */
+  propTypes?: never;
+}
+
+export type ElementTagNameMap = HTMLElementTagNameMap &
+  Pick<
+    SVGElementTagNameMap,
+    Exclude<keyof SVGElementTagNameMap, keyof HTMLElementTagNameMap>
+  >;
+
+export type ElementByTag<
+  TagName extends keyof ElementTagNameMap
+> = ElementTagNameMap[TagName];
+
 /*
 
 Test components to make sure our dynamic As prop components work as intended