diff --git a/server/router/multistage.go b/server/router/multistage.go index b28cf14..a846516 100644 --- a/server/router/multistage.go +++ b/server/router/multistage.go @@ -7,11 +7,14 @@ import ( ) // NewMultistage 创建一个支持多级分类的路由器 -func NewMultistage[HandleFunc any]() *Multistage[HandleFunc] { +func NewMultistage[HandleFunc any](options ...MultistageOption[HandleFunc]) *Multistage[HandleFunc] { r := &Multistage[HandleFunc]{ routes: make(map[any]HandleFunc), subs: make(map[any]*Multistage[HandleFunc]), } + for _, option := range options { + option(r) + } return r } @@ -28,6 +31,7 @@ type Multistage[HandleFunc any] struct { routes map[any]HandleFunc subs map[any]*Multistage[HandleFunc] tag any + trim func(route any) any } // Register 注册路由是结合 Sub 和 Route 的快捷方式,用于一次性注册多级路由 @@ -47,21 +51,25 @@ func (slf *Multistage[HandleFunc]) Register(routes ...any) MultistageBind[Handle // Route 为特定路由绑定处理函数,被绑定的处理函数将可以通过 Match 函数进行匹配 func (slf *Multistage[HandleFunc]) Route(route any, handleFunc HandleFunc) { + trim := route + if slf.trim != nil { + trim = slf.trim(route) + } if reflect.TypeOf(handleFunc).Kind() != reflect.Func { panic(fmt.Errorf("route[%v] registration failed, handle must be a function type", route)) } - _, exist := slf.routes[route] + _, exist := slf.routes[trim] if exist { panic(fmt.Errorf("the route[%v] has already been registered, duplicate registration is not allowed", route)) } - _, exist = slf.subs[route] + _, exist = slf.subs[trim] if exist { panic(fmt.Errorf("the route[%v] has already been registered, duplicate registration is not allowed", route)) } - slf.routes[route] = handleFunc + slf.routes[trim] = handleFunc if slf.tag == nil { log.Info("Router", log.String("Type", "Multistage"), log.String("Route", fmt.Sprintf("%v", route))) } else { @@ -97,12 +105,16 @@ func (slf *Multistage[HandleFunc]) Match(routes ...any) HandleFunc { // Sub 获取子路由器 func (slf *Multistage[HandleFunc]) Sub(route any) *Multistage[HandleFunc] { - _, exist := slf.routes[route] + trim := route + if slf.trim != nil { + trim = slf.trim(route) + } + _, exist := slf.routes[trim] if exist { panic(fmt.Errorf("the route[%v] has already been registered, cannot be used as a sub-router", route)) } - router, exist := slf.subs[route] + router, exist := slf.subs[trim] if !exist { router = NewMultistage[HandleFunc]() if slf.tag == nil { @@ -110,7 +122,7 @@ func (slf *Multistage[HandleFunc]) Sub(route any) *Multistage[HandleFunc] { } else { router.tag = fmt.Sprintf("%v > %v", slf.tag, route) } - slf.subs[route] = router + slf.subs[trim] = router } return router } diff --git a/server/router/multistage_options.go b/server/router/multistage_options.go new file mode 100644 index 0000000..b4bc8b7 --- /dev/null +++ b/server/router/multistage_options.go @@ -0,0 +1,12 @@ +package router + +// MultistageOption 路由器选项 +type MultistageOption[HandleFunc any] func(multistage *Multistage[HandleFunc]) + +// WithRouteTrim 路由修剪选项 +// - 将在路由注册前对路由进行对应处理 +func WithRouteTrim[HandleFunc any](handle func(route any) any) MultistageOption[HandleFunc] { + return func(multistage *Multistage[HandleFunc]) { + multistage.trim = handle + } +}