better type restriction

src/decorators/decorators.ts

@@ -9,7 +9,10 @@ import {
   KoishiCommandPutDef,
   KoishiOnContextScope,
   OnContextFunction,
+  PickEventFunction,
+  PluginDefinition,
   TopLevelActionDef,
+  TypedMethodDecorator,
 } from '../def';
 import 'reflect-metadata';
 import {
@@ -20,6 +23,8 @@ import {
   Selection,
   Session,
   I18n,
+  Awaitable,
+  Middleware,
 } from 'koishi';
 import { Metadata } from '../meta/metadata.decorators';
 import { CommandPut, DoRegister } from '../registry';
@@ -28,32 +33,43 @@ import {
   applyOptionToCommand,
   registerTemplate,
 } from '../utility';
+import { EventMap, BeforeEventMap } from 'koishi';
 
 // Register method
 
-export const UseMiddleware = (prepend = false): MethodDecorator =>
+export const UseMiddleware = (prepend = false) =>
   DoRegister.middleware(prepend);
-export const UseEvent = (name: EventName, prepend = false): MethodDecorator =>
+export const UseEvent = <K extends EventName>(
+  name: K,
+  prepend = false,
+): TypedMethodDecorator<PickEventFunction<EventMap, K>> =>
   DoRegister.event({ name, prepend });
-export const UseBeforeEvent = (
-  name: BeforeEventName,
+export const UseBeforeEvent = <K extends BeforeEventName>(
+  name: K,
   prepend = false,
-): MethodDecorator => DoRegister.beforeEvent({ name, prepend });
-export const UsePlugin = (): MethodDecorator => DoRegister.plugin();
+): TypedMethodDecorator<PickEventFunction<BeforeEventMap, K>> =>
+  DoRegister.beforeEvent({ name, prepend });
+export const UsePlugin = () => DoRegister.plugin();
 
 export function UseCommand<D extends string>(
   def: D,
   config?: CommandConfigExtended,
-): MethodDecorator;
+): TypedMethodDecorator<
+  (...args: any[]) => Awaitable<string | void | undefined>
+>;
 export function UseCommand<D extends string>(
   def: D,
   desc: string,
   config?: CommandConfigExtended,
-): MethodDecorator;
+): TypedMethodDecorator<
+  (...args: any[]) => Awaitable<string | void | undefined>
+>;
 export function UseCommand(
   def: string,
   ...args: [CommandConfigExtended?] | [string, CommandConfigExtended?]
-): MethodDecorator {
+): TypedMethodDecorator<
+  (...args: any[]) => Awaitable<string | void | undefined>
+> {
   const desc = typeof args[0] === 'string' ? (args.shift() as string) : '';
   const config = args[0] as CommandConfigExtended;
   return (obj, key: string, des) => {

src/def/interfaces.ts

@@ -7,6 +7,7 @@ import {
   Plugin,
   Selection,
   I18n,
+  Awaitable,
 } from 'koishi';
 import type { DefaultContext, DefaultState, ParameterizedContext } from 'koa';
 import type { RouterParamContext } from '@koa/router';
@@ -176,3 +177,29 @@ export interface CommandArgDef {
 }
 
 export type ParamRenderer = <T>(v: T) => T;
+
+export type TypedMethodDecorator<F extends (...args: any[]) => any> = <
+  T extends F,
+>(
+  // eslint-disable-next-line @typescript-eslint/ban-types
+  target: Object,
+  propertyKey: string | symbol,
+  descriptor: TypedPropertyDescriptor<T>,
+) => void;
+
+export type FunctionParam<F extends (...args: any[]) => any> = F extends (
+  ...args: infer R
+) => any
+  ? R
+  : never;
+export type FunctionReturn<F extends (...args: any[]) => any> = F extends (
+  ...args: any[]
+) => infer R
+  ? R
+  : never;
+
+export type PickEventFunction<M, K extends keyof M> = M[K] extends (
+  ...args: any[]
+) => any
+  ? (...args: FunctionParam<M[K]>) => Awaitable<FunctionReturn<M[K]>>
+  : M[K];

src/registry/registries/do-register.ts

 import { MethodMap, MethodRegistry } from '../abstract-registry';
-import { Argv, Awaitable, Context, MaybeArray } from 'koishi';
 import {
+  Argv,
+  Awaitable,
+  BeforeEventMap,
+  Context,
+  EventMap,
+  I18n,
+  MaybeArray,
+  Middleware,
+} from 'koishi';
+import {
+  BeforeEventName,
   BeforeEventNameAndPrepend,
   CommandRegisterConfig,
+  EventName,
   EventNameAndPrepend,
+  KoaContext,
   KoishiCommandDefinition,
   KoishiDoRegister,
   KoishiDoRegisterKeys,
   KoishiRouteDef,
   MappingStruct,
+  PickEventFunction,
   PluginDefinition,
+  TypedMethodDecorator,
 } from '../../def';
 import { Metadata } from '../../meta/metadata.decorators';
 import { reflector } from '../../meta/meta-fetch';
 import { CommandPut } from './command-put';
-import { applySelector, generateRenderer, renderObject } from '../../utility';
+import { applySelector, generateRenderer } from '../../utility';
+import { Next as KoaNext } from 'koa';
+import { IncomingMessage } from 'http';
 
 // eslint-disable-next-line @typescript-eslint/no-namespace
 export namespace DoRegister {
   export function decorate<K extends keyof ConfigMap>(
     config: Config<K>,
-  ): MethodDecorator {
+  ): K extends keyof DecoratorMap
+    ? TypedMethodDecorator<DecoratorMap[K]>
+    : MethodDecorator {
+    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+    // @ts-ignore
     return Metadata.set(KoishiDoRegister, config, KoishiDoRegisterKeys);
   }
 
@@ -44,7 +64,8 @@ export namespace DoRegister {
       method: MethodMap<ConfigMap, any, [Context, any, string, ...any[]]>[K],
     ) {
       this.extend(name, method);
-      return (config: ConfigMap[K]) => decorate({ type: name, data: config });
+      return (config: ConfigMap[K]) =>
+        decorate<K>({ type: name, data: config });
     }
   }
 
@@ -61,6 +82,18 @@ export namespace DoRegister {
     interval: number;
   }
 
+  export interface DecoratorMap {
+    middleware: Middleware;
+    onEvent: PickEventFunction<EventMap, EventName>;
+    beforeEvent: PickEventFunction<BeforeEventMap, BeforeEventName>;
+    plugin(): Awaitable<PluginDefinition | undefined>;
+    command(...args: any[]): Awaitable<string | void | undefined>;
+    route(ctx: KoaContext, Next: KoaNext): any;
+    ws(socket: WebSocket, request: IncomingMessage): any;
+    formatter: I18n.Formatter;
+    preset: I18n.Renderer;
+  }
+
   export type Config<K extends keyof ConfigMap = keyof ConfigMap> =
     MappingStruct<ConfigMap, K>;
 

tests/register.spec.ts

@@ -26,13 +26,13 @@ class MyClass {
     if (session.content === 'ping') {
       return 'pong';
     }
-    return next;
+    return next();
   }
 
   @UseEvent('message')
   async onMessage(session: Session) {
     if (session.content === 'pang') {
-      return session.send('peng');
+      await session.send('peng');
     }
   }