Merge pull request #52 from sqs/BuildVarsFunc

BuildVarsFunc: modify route variables before building the URL from a route

mux.go

@@ -152,6 +152,13 @@ func (r *Router) getRegexpGroup() *routeRegexpGroup {
 	return nil
 }
 
+func (r *Router) buildVars(m map[string]string) map[string]string {
+	if r.parent != nil {
+		m = r.parent.buildVars(m)
+	}
+	return m
+}
+
 // ----------------------------------------------------------------------------
 // Route factories
 // ----------------------------------------------------------------------------
@@ -224,6 +231,12 @@ func (r *Router) Schemes(schemes ...string) *Route {
 	return r.NewRoute().Schemes(schemes...)
 }
 
+// BuildVars registers a new route with a custom function for modifying
+// route variables before building a URL.
+func (r *Router) BuildVarsFunc(f BuildVarsFunc) *Route {
+	return r.NewRoute().BuildVarsFunc(f)
+}
+
 // ----------------------------------------------------------------------------
 // Context
 // ----------------------------------------------------------------------------

mux_test.go

@@ -593,6 +593,39 @@ func TestMatcherFunc(t *testing.T) {
 	}
 }
 
+func TestBuildVarsFunc(t *testing.T) {
+	tests := []routeTest{
+		{
+			title: "BuildVarsFunc set on route",
+			route: new(Route).Path(`/111/{v1:\d}{v2:.*}`).BuildVarsFunc(func(vars map[string]string) map[string]string {
+				vars["v1"] = "3"
+				vars["v2"] = "a"
+				return vars
+			}),
+			request:     newRequest("GET", "http://localhost/111/2"),
+			path:        "/111/3a",
+			shouldMatch: true,
+		},
+		{
+			title: "BuildVarsFunc set on route and parent route",
+			route: new(Route).PathPrefix(`/{v1:\d}`).BuildVarsFunc(func(vars map[string]string) map[string]string {
+				vars["v1"] = "2"
+				return vars
+			}).Subrouter().Path(`/{v2:\w}`).BuildVarsFunc(func(vars map[string]string) map[string]string {
+				vars["v2"] = "b"
+				return vars
+			}),
+			request:     newRequest("GET", "http://localhost/1/a"),
+			path:        "/2/b",
+			shouldMatch: true,
+		},
+	}
+
+	for _, test := range tests {
+		testRoute(t, test)
+	}
+}
+
 func TestSubRouter(t *testing.T) {
 	subrouter1 := new(Route).Host("{v1:[a-z]+}.google.com").Subrouter()
 	subrouter2 := new(Route).PathPrefix("/foo/{v1}").Subrouter()

regexp.go

@@ -150,11 +150,7 @@ func (r *routeRegexp) Match(req *http.Request, match *RouteMatch) bool {
 }
 
 // url builds a URL part using the given values.
-func (r *routeRegexp) url(pairs ...string) (string, error) {
+func (r *routeRegexp) url(values map[string]string) (string, error) {
-	values, err := mapFromPairs(pairs...)
-	if err != nil {
-		return "", err
-	}
 	urlValues := make([]interface{}, len(r.varsN))
 	for k, v := range r.varsN {
 		value, ok := values[v]

route.go

@@ -31,6 +31,8 @@ type Route struct {
 	name string
 	// Error resulted from building a route.
 	err error
+
+	buildVarsFunc BuildVarsFunc
 }
 
 // Match matches the route against the request.
@@ -360,6 +362,19 @@ func (r *Route) Schemes(schemes ...string) *Route {
 	return r.addMatcher(schemeMatcher(schemes))
 }
 
+// BuildVarsFunc --------------------------------------------------------------
+
+// BuildVarsFunc is the function signature used by custom build variable
+// functions (which can modify route variables before a route's URL is built).
+type BuildVarsFunc func(map[string]string) map[string]string
+
+// BuildVarsFunc adds a custom function to be used to modify build variables
+// before a route's URL is built.
+func (r *Route) BuildVarsFunc(f BuildVarsFunc) *Route {
+	r.buildVarsFunc = f
+	return r
+}
+
 // Subrouter ------------------------------------------------------------------
 
 // Subrouter creates a subrouter for the route.
@@ -422,17 +437,20 @@ func (r *Route) URL(pairs ...string) (*url.URL, error) {
 	if r.regexp == nil {
 		return nil, errors.New("mux: route doesn't have a host or path")
 	}
+	values, err := r.prepareVars(pairs...)
+	if err != nil {
+		return nil, err
+	}
 	var scheme, host, path string
-	var err error
 	if r.regexp.host != nil {
 		// Set a default scheme.
 		scheme = "http"
-		if host, err = r.regexp.host.url(pairs...); err != nil {
+		if host, err = r.regexp.host.url(values); err != nil {
 			return nil, err
 		}
 	}
 	if r.regexp.path != nil {
-		if path, err = r.regexp.path.url(pairs...); err != nil {
+		if path, err = r.regexp.path.url(values); err != nil {
 			return nil, err
 		}
 	}
@@ -453,7 +471,11 @@ func (r *Route) URLHost(pairs ...string) (*url.URL, error) {
 	if r.regexp == nil || r.regexp.host == nil {
 		return nil, errors.New("mux: route doesn't have a host")
 	}
-	host, err := r.regexp.host.url(pairs...)
+	values, err := r.prepareVars(pairs...)
+	if err != nil {
+		return nil, err
+	}
+	host, err := r.regexp.host.url(values)
 	if err != nil {
 		return nil, err
 	}
@@ -473,7 +495,11 @@ func (r *Route) URLPath(pairs ...string) (*url.URL, error) {
 	if r.regexp == nil || r.regexp.path == nil {
 		return nil, errors.New("mux: route doesn't have a path")
 	}
-	path, err := r.regexp.path.url(pairs...)
+	values, err := r.prepareVars(pairs...)
+	if err != nil {
+		return nil, err
+	}
+	path, err := r.regexp.path.url(values)
 	if err != nil {
 		return nil, err
 	}
@@ -482,6 +508,26 @@ func (r *Route) URLPath(pairs ...string) (*url.URL, error) {
 	}, nil
 }
 
+// prepareVars converts the route variable pairs into a map. If the route has a
+// BuildVarsFunc, it is invoked.
+func (r *Route) prepareVars(pairs ...string) (map[string]string, error) {
+	m, err := mapFromPairs(pairs...)
+	if err != nil {
+		return nil, err
+	}
+	return r.buildVars(m), nil
+}
+
+func (r *Route) buildVars(m map[string]string) map[string]string {
+	if r.parent != nil {
+		m = r.parent.buildVars(m)
+	}
+	if r.buildVarsFunc != nil {
+		m = r.buildVarsFunc(m)
+	}
+	return m
+}
+
 // ----------------------------------------------------------------------------
 // parentRoute
 // ----------------------------------------------------------------------------
@@ -490,6 +536,7 @@ func (r *Route) URLPath(pairs ...string) (*url.URL, error) {
 type parentRoute interface {
 	getNamedRoutes() map[string]*Route
 	getRegexpGroup() *routeRegexpGroup
+	buildVars(map[string]string) map[string]string
 }
 
 // getNamedRoutes returns the map where named routes are registered.