From a261d98518986a4bef516bd99d2520ecc17f2db7 Mon Sep 17 00:00:00 2001
From: Yanyan-Wang <yanyanwang93@gmail.com>
Date: Mon, 16 Mar 2020 16:58:10 +0800
Subject: [PATCH] fix: behavior type error. closes: #1303.

---
 src/behavior/behavior.ts | 28 ++++++++++++++++++++++++++--
 src/types/index.ts       |  1 -
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/src/behavior/behavior.ts b/src/behavior/behavior.ts
index 77fe965017..6e156004f3 100644
--- a/src/behavior/behavior.ts
+++ b/src/behavior/behavior.ts
@@ -1,6 +1,30 @@
 import { clone, each, wrapBehavior } from '@antv/util/lib';
-import { BehaviorOption } from '../types';
+import { ModelConfig, DefaultBehaviorType } from '../types';
 import behaviorOption from './behaviorOption';
+import { IGraph } from '../interface/graph';
+import { G6Event } from '../types';
+
+type GetEvents = 'getEvents';
+type ShouldBegin = 'shouldBegin';
+type ShouldUpdate = 'shouldUpdate';
+type ShouldEnd = 'shouldEnd';
+type Bind = 'bind';
+type Unbind = 'unbind';
+type BehaviorOption<U> = {
+  [T in keyof U]: T extends GetEvents
+    ? () => { [key in G6Event]?: string }
+    : T extends ShouldBegin
+    ? (cfg?: ModelConfig) => boolean
+    : T extends ShouldEnd
+    ? (cfg?: ModelConfig) => boolean
+    : T extends ShouldUpdate
+    ? (cfg?: ModelConfig) => boolean
+    : T extends Bind
+    ? (graph: IGraph) => void
+    : T extends Unbind
+    ? (graph: IGraph) => void
+    : (...args: DefaultBehaviorType[]) => unknown;
+};
 
 export default class Behavior {
   // 所有自定义的 Behavior 的实例
@@ -11,7 +35,7 @@ export default class Behavior {
    * @param type Behavior 名称
    * @param behavior Behavior 定义的方法集合
    */
-  public static registerBehavior(type: string, behavior: BehaviorOption) {
+  public static registerBehavior<T, U>(type: string, behavior: BehaviorOption<U>) {
     if (!behavior) {
       throw new Error(`please specify handler for this behavior: ${type}`);
     }
diff --git a/src/types/index.ts b/src/types/index.ts
index 5a75ad7660..07f30745c6 100644
--- a/src/types/index.ts
+++ b/src/types/index.ts
@@ -352,7 +352,6 @@ export interface BehaviorOption {
   shouldEnd?(e?: IG6GraphEvent): boolean;
   bind?(e: IGraph): void;
   unbind?(e: IGraph): void;
-  [key: string]: (...args) => any;
 }
 
 export type IEvent = Record<G6Event, string>;