From 1ac3032b21f98592663549c1c27d70efcdff9dfd Mon Sep 17 00:00:00 2001
From: Max Davidson <max.davidson@cygni.se>
Date: Thu, 27 Sep 2018 13:09:49 +0200
Subject: [PATCH] Improve HOC type inference

---
 types/react-navigation/index.d.ts             | 36 ++++++++++++++++---
 .../react-navigation-tests.tsx                | 32 +++++++++++------
 2 files changed, 53 insertions(+), 15 deletions(-)

diff --git a/types/react-navigation/index.d.ts b/types/react-navigation/index.d.ts
index e671e6c9ac..80b54d23a2 100644
--- a/types/react-navigation/index.d.ts
+++ b/types/react-navigation/index.d.ts
@@ -24,6 +24,7 @@
 //                 Gustavo Brunoro <https://github.com/brunoro>
 //                 Denis Frezzato <https://github.com/DenisFrezzato>
 //                 Mickael Wegerich <https://github.com/mickaelw>
+//                 Max Davidson <https://github.com/maxdavidson>
 // Definitions: https://github.com/DefinitelyTyped/DefinitelyTyped
 // TypeScript Version: 2.8
 
@@ -1234,20 +1235,45 @@ export class Header extends React.Component<HeaderProps> {
   static HEIGHT: number;
 }
 
+export type Omit<T, K extends keyof any> = Pick<T, Exclude<keyof T, K>>;
+
+export type InferProps<T extends React.ComponentType<any>> = T extends React.ComponentType<infer P> ? P : never;
+
 export interface NavigationInjectedProps<P = NavigationParams> {
   navigation: NavigationScreenProp<NavigationState, P>;
 }
 
+// If the wrapped component is a class, we can get a ref to it
+export function withNavigation<T extends React.ComponentClass<NavigationInjectedProps>>(
+  Component: T,
+): React.ComponentType<Omit<InferProps<T>, keyof NavigationInjectedProps> & { onRef?: React.Ref<InstanceType<T>> }>;
+
+export function withNavigation<T extends React.ComponentType<NavigationInjectedProps>>(
+  Component: T,
+): React.ComponentType<Omit<InferProps<T>, keyof NavigationInjectedProps>>;
+
+// For backwards compatibility
 export function withNavigation<T = {}, P = NavigationParams>(
-  Component: React.ComponentType<T & NavigationInjectedProps<P>>
+  Component: React.ComponentType<T & NavigationInjectedProps<P>>,
 ): React.ComponentType<T & { onRef?: React.Ref<React.Component<T & NavigationInjectedProps<P>>> }>;
 
-export interface NavigationFocusInjectedProps extends NavigationInjectedProps {
+export interface NavigationFocusInjectedProps<P = NavigationParams> extends NavigationInjectedProps<P> {
   isFocused: boolean;
 }
-export function withNavigationFocus<T = {}>(
-  Component: React.ComponentType<T & NavigationFocusInjectedProps>
-): React.ComponentType<T & { onRef?: React.Ref<typeof Component> }>;
+
+// If the wrapped component is a class, we can get a ref to it
+export function withNavigationFocus<T extends React.ComponentClass<NavigationFocusInjectedProps>>(
+  Component: T,
+): React.ComponentType<Omit<InferProps<T>, keyof NavigationFocusInjectedProps> & { onRef?: React.Ref<InstanceType<T>> }>;
+
+export function withNavigationFocus<T extends React.ComponentType<NavigationFocusInjectedProps>>(
+  Component: T,
+): React.ComponentType<Omit<InferProps<T>, keyof NavigationFocusInjectedProps>>;
+
+// For backwards compatibility
+export function withNavigationFocus<T = {}, P = NavigationParams>(
+  Component: React.ComponentType<T & NavigationFocusInjectedProps<P>>,
+): React.ComponentType<T & { onRef?: React.Ref<React.Component<T & NavigationFocusInjectedProps<P>>> }>;
 
 /**
  * SafeAreaView Component
diff --git a/types/react-navigation/react-navigation-tests.tsx b/types/react-navigation/react-navigation-tests.tsx
index a9063577a6..47dd7baa0e 100644
--- a/types/react-navigation/react-navigation-tests.tsx
+++ b/types/react-navigation/react-navigation-tests.tsx
@@ -550,20 +550,32 @@ class MyBackButton extends React.Component<BackButtonProps & NavigationInjectedP
     }
 }
 
-// withNavigation returns a component that wraps MyBackButton and passes in the
-// navigation prop
-const BackButtonWithNavigation = withNavigation<BackButtonProps>(MyBackButton);
+// withNavigation returns a component that wraps MyBackButton and passes in the navigation prop.
+// If you have class methods, you should have a way to use them.
+const BackButtonWithNavigation = withNavigation(MyBackButton);
 const BackButtonInstance = <BackButtonWithNavigation
-    title="Back" onRef={ref => { const backButtonRef = ref; }}
+    title="Back" onRef={ref => {
+        // ref is inferred as MyBackButton | null
+        if (!ref) return;
+        ref.triggerBack();
+    }}
 />;
 
-// if you have class methods, you should have a way to use them
-const BackButtonWithNavigationSpecified = withNavigation<BackButtonProps>(MyBackButton);
-const BackButtonSpecifiedInstance = <BackButtonWithNavigationSpecified
+function StatelessBackButton(props: BackButtonProps & NavigationInjectedProps) {
+  return <MyBackButton {...props} />;
+}
+
+// Wrapped stateless components don't accept an onRef
+const StatelessBackButtonWithNavigation = withNavigation(StatelessBackButton);
+const StatelessBackButtonInstance = <StatelessBackButtonWithNavigation title="Back" />;
+
+// The old way of passing in the props should still work
+const BackButtonWithNavigationWithExplicitProps = withNavigation<BackButtonProps>(MyBackButton);
+const BackButtonWithExplicitPropsInstance = <BackButtonWithNavigationWithExplicitProps
     title="Back" onRef={ref => {
         if (!ref) return;
-        const backButtonRef = ref as MyBackButton;
-        backButtonRef.triggerBack();
+        // We can't infer the component type if we pass in the props
+        (ref as MyBackButton).triggerBack();
     }}
 />;
 
@@ -578,7 +590,7 @@ class MyFocusedComponent extends React.Component<MyFocusedComponentProps & Navig
 
 // withNavigationFocus returns a component that wraps MyFocusedComponent and passes in the
 // navigation and isFocused prop
-const MyFocusedComponentWithNavigationFocus = withNavigationFocus<MyFocusedComponentProps>(MyFocusedComponent);
+const MyFocusedComponentWithNavigationFocus = withNavigationFocus(MyFocusedComponent);
 const MyFocusedComponentInstance = <MyFocusedComponentWithNavigationFocus
     expectsFocus={true} onRef={ref => { const backButtonRef = ref; }}
 />;