stronger withModel

package.json

@@ -43,5 +43,7 @@
     "tslint": "^5.18.0",
     "typescript": "^3.4.5"
   },
-  "dependencies": {}
+  "dependencies": {
+    "@types/hoist-non-react-statics": "^3.3.1"
+  }
 }

src/__tests__/__snapshots__/main.test.tsx.snap

@@ -21,3 +21,25 @@ exports[`provider initialize 2`] = `
   </div>
 </DocumentFragment>
 `;
+
+exports[`withModel 1`] = `
+<DocumentFragment>
+  <div>
+    <button>
+      Change
+    </button>
+    0
+  </div>
+</DocumentFragment>
+`;
+
+exports[`withModel 2`] = `
+<DocumentFragment>
+  <div>
+    <button>
+      Change
+    </button>
+    1
+  </div>
+</DocumentFragment>
+`;

src/__tests__/main.test.tsx

@@ -1,6 +1,6 @@
-import { setModel, useModel, selectModel } from "..";
+import { setModel, useModel, selectModel, withModel } from "..";
 import * as React from "react";
-import { FC, memo, useState } from "react";
+import { Component, FC, memo, useState } from "react";
 import * as testing from "@testing-library/react";
 
 test("provider initialize", function() {
@@ -28,3 +28,39 @@ test("provider initialize", function() {
   testing.fireEvent.click(testing.getByText(renderer.container, "Change"));
   expect(renderer.asFragment()).toMatchSnapshot();
 });
+
+test("withModel", function() {
+  function useCounter() {
+    const [count, setCount] = useState(0);
+    const decrement = () => setCount(count - 1);
+    const increment = () => setCount(count + 1);
+    return { count, decrement, increment };
+  }
+
+  setModel("counter", useCounter);
+
+  interface Props {
+    counter: ReturnType<typeof useCounter>;
+  }
+
+  class App extends Component<Props> {
+    render() {
+      return (
+        <div>
+          <button onClick={this.props.counter.increment}>Change</button>
+          {this.props.counter.count}
+        </div>
+      );
+    }
+  }
+
+  type Counter = ReturnType<typeof useCounter>;
+
+  const AppWithModel = withModel("counter", (model: { counter: Counter }) => ({
+    counter: model.counter
+  }))(App);
+  const renderer = testing.render(<AppWithModel />);
+  expect(renderer.asFragment()).toMatchSnapshot();
+  testing.fireEvent.click(testing.getByText(renderer.container, "Change"));
+  expect(renderer.asFragment()).toMatchSnapshot();
+});

src/index.tsx

@@ -6,8 +6,10 @@ import React, {
   useEffect,
   useRef,
   useState,
-  ReactElement
+  ReactElement,
+  NamedExoticComponent
 } from "react";
+import { NonReactStatics } from "hoist-non-react-statics";
 
 type ModelHook<T = any> = () => T;
 
@@ -96,26 +98,91 @@ export function selectModel<T extends ModelHook = any>(key: string) {
   return container ? (container.data as V) : undefined;
 }
 
+export type Omit<T, K extends keyof T> = Pick<T, Exclude<keyof T, K>>;
+
+export type GetProps<C> = C extends ComponentType<infer P> ? P : never;
+
+export type ConnectedComponent<
+  C extends ComponentType<any>,
+  P
+> = NamedExoticComponent<JSX.LibraryManagedAttributes<C, P>> &
+  NonReactStatics<C> & {
+    WrappedComponent: C;
+  };
+
+export type Matching<InjectedProps, DecorationTargetProps> = {
+  [P in keyof DecorationTargetProps]: P extends keyof InjectedProps
+    ? InjectedProps[P] extends DecorationTargetProps[P]
+      ? DecorationTargetProps[P]
+      : InjectedProps[P]
+    : DecorationTargetProps[P];
+};
+
+export type Shared<InjectedProps, DecorationTargetProps> = {
+  [P in Extract<
+    keyof InjectedProps,
+    keyof DecorationTargetProps
+  >]?: InjectedProps[P] extends DecorationTargetProps[P]
+    ? DecorationTargetProps[P]
+    : never;
+};
+
+export type InferableComponentEnhancerWithProps<TInjectedProps, TNeedsProps> = <
+  C extends ComponentType<Matching<TInjectedProps, GetProps<C>>>
+>(
+  component: C
+) => ConnectedComponent<
+  C,
+  Omit<GetProps<C>, keyof Shared<TInjectedProps, GetProps<C>>> & TNeedsProps
+>;
+
 export interface WithModelProps {
   model: {
     [key: string]: unknown;
   };
 }
 
-export function withModel<T extends ModelHook = any>(key: string) {
-  return function<P extends WithModelProps>(C: ComponentType<P>) {
-    const Wrapper: FC<Omit<P, "model">> = function(props) {
-      const componentProps: P = ({
-        ...props,
-        model: {
-          [key]: useModel<T>(key)
+type MapModelToProps<TModelProps, TOwnProps, Model> = (
+  model: Model,
+  ownProps: TOwnProps
+) => TModelProps;
+
+export function withModel<TModelProps, TOwnProps, Model>(
+  key: string,
+  mapModelToProps?: MapModelToProps<TModelProps, TOwnProps, Model>
+): InferableComponentEnhancerWithProps<TModelProps, TOwnProps>;
+export function withModel<TModelProps, TOwnProps, Model>(
+  keys: string[],
+  mapModelToProps?: MapModelToProps<TModelProps, TOwnProps, Model>
+): InferableComponentEnhancerWithProps<TModelProps, TOwnProps>;
+export function withModel<TModelProps, TOwnProps, Model>(
+  keyOrKeys: string | string[],
+  mapModelToProps: MapModelToProps<TModelProps, TOwnProps, Model>
+) {
+  return function(C) {
+    const Wrapper: FC<any> = function(props) {
+      const model: {
+        [key: string]: unknown;
+      } = {};
+      if (Array.isArray(keyOrKeys)) {
+        for (const key of keyOrKeys) {
+          model[key] = useModel(key);
         }
-      } as unknown) as P;
+      } else {
+        model[keyOrKeys] = useModel(keyOrKeys);
+      }
+      const modelProps = mapModelToProps
+        ? mapModelToProps((model as unknown) as Model, props)
+        : { model };
+      const componentProps = {
+        ...props,
+        ...modelProps
+      };
       return <C {...componentProps} />;
     };
     Wrapper.displayName = `${C.displayName}Wrapper`;
     return Wrapper;
-  };
+  } as InferableComponentEnhancerWithProps<TModelProps, TOwnProps>;
 }
 
 function compare(oldDeps: unknown[], newDeps: unknown[]) {

yarn.lock

@@ -407,6 +407,14 @@
     "@types/minimatch" "*"
     "@types/node" "*"
 
+"@types/hoist-non-react-statics@^3.3.1":
+  version "3.3.1"
+  resolved "https://registry.npm.taobao.org/@types/hoist-non-react-statics/download/@types/hoist-non-react-statics-3.3.1.tgz#1124aafe5118cb591977aeb1ceaaed1070eb039f"
+  integrity sha1-ESSq/lEYy1kZd66xzqrtEHDrA58=
+  dependencies:
+    "@types/react" "*"
+    hoist-non-react-statics "^3.3.0"
+
 "@types/istanbul-lib-coverage@*", "@types/istanbul-lib-coverage@^2.0.0":
   version "2.0.1"
   resolved "https://registry.npm.taobao.org/@types/istanbul-lib-coverage/download/@types/istanbul-lib-coverage-2.0.1.tgz#42995b446db9a48a11a07ec083499a860e9138ff"
@@ -2246,6 +2254,13 @@ has@^1.0.1, has@^1.0.3:
   dependencies:
     function-bind "^1.1.1"
 
+hoist-non-react-statics@^3.3.0:
+  version "3.3.0"
+  resolved "https://registry.npm.taobao.org/hoist-non-react-statics/download/hoist-non-react-statics-3.3.0.tgz#b09178f0122184fb95acf525daaecb4d8f45958b"
+  integrity sha1-sJF48BIhhPuVrPUl2q7LTY9FlYs=
+  dependencies:
+    react-is "^16.7.0"
+
 homedir-polyfill@^1.0.1:
   version "1.0.3"
   resolved "https://registry.npm.taobao.org/homedir-polyfill/download/homedir-polyfill-1.0.3.tgz#743298cef4e5af3e194161fbadcc2151d3a058e8"
@@ -4209,7 +4224,7 @@ react-dom@^16.8.0:
     prop-types "^15.6.2"
     scheduler "^0.16.2"
 
-react-is@^16.8.1, react-is@^16.8.4, react-is@^16.8.6:
+react-is@^16.7.0, react-is@^16.8.1, react-is@^16.8.4, react-is@^16.8.6:
   version "16.10.2"
   resolved "https://registry.npm.taobao.org/react-is/download/react-is-16.10.2.tgz#984120fd4d16800e9a738208ab1fba422d23b5ab"
   integrity sha1-mEEg/U0WgA6ac4IIqx+6Qi0jtas=