beego:addtree support regexp

tree.go

@@ -28,21 +28,96 @@ func NewTree() *Tree {
 // add Tree to the exist Tree
 // prefix should has no params
 func (t *Tree) AddTree(prefix string, tree *Tree) {
-	t.addtree(splitPath(prefix), tree)
+	t.addtree(splitPath(prefix), tree, nil, "")
 }
 
-func (t *Tree) addtree(segments []string, tree *Tree) {
+func (t *Tree) addtree(segments []string, tree *Tree, wildcards []string, reg string) {
 	if len(segments) == 0 {
 		panic("prefix should has path")
 	}
-	if len(segments) == 1 && segments[0] != "" {
-		t.fixrouters[segments[0]] = tree
+	seg := segments[0]
+	iswild, params, regexpStr := splitSegment(seg)
+	if len(segments) == 1 && seg != "" {
+		if iswild {
+			wildcards = append(wildcards, params...)
+			if regexpStr != "" {
+				for _, w := range params {
+					if w == "." || w == ":" {
+						continue
+					}
+					regexpStr = "([^/]+)/" + regexpStr
+				}
+			}
+			reg = reg + regexpStr
+			filterTreeWithPrefix(tree, wildcards, reg)
+			t.wildcard = tree
+		} else {
+			filterTreeWithPrefix(tree, wildcards, reg)
+			t.fixrouters[seg] = tree
+		}
 		return
 	}
-	seg := segments[0]
-	subTree := NewTree()
-	t.fixrouters[seg] = subTree
-	subTree.addtree(segments[1:], tree)
+	if iswild {
+		if t.wildcard == nil {
+			t.wildcard = NewTree()
+		}
+		wildcards = append(wildcards)
+		if regexpStr != "" {
+			for _, w := range params {
+				if w == "." || w == ":" {
+					continue
+				}
+				regexpStr = "([^/]+)/" + regexpStr
+			}
+		}
+		reg = reg + regexpStr
+		t.wildcard.addtree(segments[1:], tree, wildcards, reg)
+	} else {
+		subTree := NewTree()
+		t.fixrouters[seg] = subTree
+		subTree.addtree(segments[1:], tree, wildcards, reg)
+	}
+}
+
+func filterTreeWithPrefix(t *Tree, wildcards []string, reg string) {
+	for _, v := range t.fixrouters {
+		filterTreeWithPrefix(v, wildcards, reg)
+	}
+	if t.wildcard != nil {
+		filterTreeWithPrefix(t.wildcard, wildcards, reg)
+	}
+	if t.leaf != nil {
+		t.leaf.wildcards = append(wildcards, t.leaf.wildcards...)
+		if reg != "" {
+			filterCards := []string{}
+			for _, v := range t.leaf.wildcards {
+				if v == ":" || v == "." {
+					continue
+				}
+				filterCards = append(filterCards, v)
+			}
+			t.leaf.wildcards = filterCards
+			t.leaf.regexps = regexp.MustCompile("^" + reg + strings.Trim(t.leaf.regexps.String(), "^$") + "$")
+		} else {
+			if t.leaf.regexps != nil {
+				filterCards := []string{}
+				for _, v := range t.leaf.wildcards {
+					if v == ":" || v == "." {
+						continue
+					}
+					filterCards = append(filterCards, v)
+				}
+				t.leaf.wildcards = filterCards
+				for _, w := range wildcards {
+					if w == "." || w == ":" {
+						continue
+					}
+					reg = "([^/]+)/" + reg
+				}
+				t.leaf.regexps = regexp.MustCompile("^" + reg + strings.Trim(t.leaf.regexps.String(), "^$") + "$")
+			}
+		}
+	}
 }
 
 // call addseg function
@@ -74,6 +149,18 @@ func (t *Tree) addseg(segments []string, route interface{}, wildcards []string,
 			if t.wildcard == nil {
 				t.wildcard = NewTree()
 			}
+			if regexpStr != "" {
+				if reg == "" {
+					for _, w := range wildcards {
+						if w == "." || w == ":" {
+							continue
+						}
+						regexpStr = "([^/]+)/" + regexpStr
+					}
+				} else {
+					regexpStr = "/" + regexpStr
+				}
+			}
 			t.wildcard.addseg(segments[1:], route, append(wildcards, params...), reg+regexpStr)
 		} else {
 			subTree, ok := t.fixrouters[seg]
@@ -230,12 +317,11 @@ func (leaf *leafInfo) match(wildcardValues []string) (ok bool, params map[string
 		}
 		return true, params
 	}
-
-	if !leaf.regexps.MatchString(strings.Join(wildcardValues, "")) {
+	if !leaf.regexps.MatchString(path.Join(wildcardValues...)) {
 		return false, nil
 	}
 	params = make(map[string]string)
-	matches := leaf.regexps.FindStringSubmatch(strings.Join(wildcardValues, ""))
+	matches := leaf.regexps.FindStringSubmatch(path.Join(wildcardValues...))
 	for i, match := range matches[1:] {
 		params[leaf.wildcards[i]] = match
 	}

tree_test.go

@@ -31,6 +31,9 @@ func init() {
 	routers = append(routers, testinfo{"/v1/shop/:id([0-9]+)_:name", "/v1/shop/123_nike", map[string]string{":id": "123", ":name": "nike"}})
 	routers = append(routers, testinfo{"/v1/shop/:id(.+)_cms.html", "/v1/shop/123_cms.html", map[string]string{":id": "123"}})
 	routers = append(routers, testinfo{"/v1/shop/cms_:id(.+)_:page(.+).html", "/v1/shop/cms_123_1.html", map[string]string{":id": "123", ":page": "1"}})
+	routers = append(routers, testinfo{"/v1/:v/cms/aaa_:id(.+)_:page(.+).html", "/v1/2/cms/aaa_123_1.html", map[string]string{":v": "2", ":id": "123", ":page": "1"}})
+	routers = append(routers, testinfo{"/v1/:v/cms_:id(.+)_:page(.+).html", "/v1/2/cms_123_1.html", map[string]string{":v": "2", ":id": "123", ":page": "1"}})
+	routers = append(routers, testinfo{"/v1/:v(.+)_cms/ttt_:id(.+)_:page(.+).html", "/v1/2_cms/ttt_123_1.html", map[string]string{":v": "2", ":id": "123", ":page": "1"}})
 }
 
 func TestTreeRouters(t *testing.T) {
@@ -53,6 +56,57 @@ func TestTreeRouters(t *testing.T) {
 	}
 }
 
+func TestAddTree(t *testing.T) {
+	tr := NewTree()
+	tr.AddRouter("/shop/:id/account", "astaxie")
+	tr.AddRouter("/shop/:sd/ttt_:id(.+)_:page(.+).html", "astaxie")
+	t1 := NewTree()
+	t1.AddTree("/v1/zl", tr)
+	obj, param := t1.Match("/v1/zl/shop/123/account")
+	if obj == nil || obj.(string) != "astaxie" {
+		t.Fatal("/v1/zl/shop/:id/account can't get obj ")
+	}
+	if param == nil {
+		t.Fatal("get param error")
+	}
+	if param[":id"] != "123" {
+		t.Fatal("get :id param error")
+	}
+	obj, param = t1.Match("/v1/zl/shop/123/ttt_1_12.html")
+	if obj == nil || obj.(string) != "astaxie" {
+		t.Fatal("/v1/zl//shop/:sd/ttt_:id(.+)_:page(.+).html can't get obj ")
+	}
+	if param == nil {
+		t.Fatal("get param error")
+	}
+	if param[":sd"] != "123" || param[":id"] != "1" || param[":page"] != "12" {
+		t.Fatal("get :sd :id :page param error")
+	}
+
+	t2 := NewTree()
+	t2.AddTree("/v1/:shopid", tr)
+	obj, param = t2.Match("/v1/zl/shop/123/account")
+	if obj == nil || obj.(string) != "astaxie" {
+		t.Fatal("/v1/:shopid/shop/:id/account can't get obj ")
+	}
+	if param == nil {
+		t.Fatal("get param error")
+	}
+	if param[":id"] != "123" || param[":shopid"] != "zl" {
+		t.Fatal("get :id :shopid param error")
+	}
+	obj, param = t2.Match("/v1/zl/shop/123/ttt_1_12.html")
+	if obj == nil || obj.(string) != "astaxie" {
+		t.Fatal("/v1/:shopid/shop/:sd/ttt_:id(.+)_:page(.+).html can't get obj ")
+	}
+	if param == nil {
+		t.Fatal("get :shopid param error")
+	}
+	if param[":sd"] != "123" || param[":id"] != "1" || param[":page"] != "12" || param[":shopid"] != "zl" {
+		t.Fatal("get :sd :id :page :shopid param error")
+	}
+}
+
 func TestSplitPath(t *testing.T) {
 	a := splitPath("/")
 	if len(a) != 0 {