Add support for import alias

README.md

@@ -100,6 +100,7 @@ func IgnoreHandler(w http.ResponseWriter, req *http.Request) {
   - The function/method signature has `*http.Request`.
       - `defer newrelic.FromContext(req.Context()).StartSegment("func_name").End()`
 - [x] Support any variable name of `context.Context`/`*http.Request`.
+- [x] Support import alias of `"context"`/`"net/http"`.
 - [x] Use function/method name to segment name.
 - [x] This processing is recursively repeated.
 - [x] Able to ignore function/method by `nrseg:ignore` comment.

inspect.go

@@ -33,7 +33,7 @@ func (nrseg *nrseg) Inspect(filename string, src []byte) error {
 				return false
 			}
 			if fd.Body != nil && len(fd.Body.List) > 0 {
-				if _, t := parseParams(fd.Type); !(t == TypeContext || t == TypeHttpRequest) {
+				if _, t := parseParams(f.Imports, fd.Type); !(t == TypeContext || t == TypeHttpRequest) {
 					return false
 				}
 				if !existFromContext(pkg, fd.Body.List[0]) {

process.go

@@ -42,7 +42,7 @@ func Process(filename string, src []byte) ([]byte, error) {
 			}
 			if fd.Body != nil && len(fd.Body.List) > 0 {
 				sn := getSegName(fd)
-				vn, t := parseParams(fd.Type)
+				vn, t := parseParams(f.Imports, fd.Type)
 				var ds ast.Stmt
 				switch t {
 				case TypeContext:
@@ -225,19 +225,24 @@ const (
 	TypeUnknown     = "Unknown"
 )
 
-func parseParams(t *ast.FuncType) (string, string) {
+var types = map[string]string{
+	TypeContext:     "\"context\"",
+	TypeHttpRequest: "\"net/http\"",
+}
+
+func parseParams(is []*ast.ImportSpec, t *ast.FuncType) (string, string) {
+	var cname = getImportName(is, TypeContext)
+	var hname = getImportName(is, TypeHttpRequest)
 	n, typ := "", TypeUnknown
 	for _, f := range t.Params.List {
 		if se, ok := f.Type.(*ast.SelectorExpr); ok {
-			// TODO: support named import
-			if idt, ok := se.X.(*ast.Ident); ok && idt.Name == "context" && se.Sel.Name == "Context" {
+			if idt, ok := se.X.(*ast.Ident); ok && idt.Name == cname && se.Sel.Name == "Context" {
 				return f.Names[0].Name, TypeContext
 			}
 		}
 		if se, ok := f.Type.(*ast.StarExpr); ok {
 			if se, ok := se.X.(*ast.SelectorExpr); ok {
-				// TODO: support named import
-				if idt, ok := se.X.(*ast.Ident); ok && idt.Name == "http" && se.Sel.Name == "Request" {
+				if idt, ok := se.X.(*ast.Ident); ok && idt.Name == hname && se.Sel.Name == "Request" {
 					n = f.Names[0].Name
 					typ = TypeHttpRequest
 				}
@@ -246,3 +251,13 @@ func parseParams(t *ast.FuncType) (string, string) {
 	}
 	return n, typ
 }
+
+func getImportName(is []*ast.ImportSpec, typ string) string {
+	var def = strings.Replace(strings.Split(typ, ".")[0], "*", "", 1)
+	for _, i := range is {
+		if i.Name != nil && i.Path != nil && i.Path.Value == types[typ] {
+			return i.Name.Name
+		}
+	}
+	return def
+}

process_test.go

@@ -331,6 +331,19 @@ import (
 )
 
 func Hoge(ctx context.Context) {}
+`,
+			wantName: "ctx", wantType: TypeContext,
+		},
+		{
+			name: "NamedContext",
+			src: `
+package main
+
+import (
+	c "context"
+)
+
+func Hoge(ctx c.Context) {}
 `,
 			wantName: "ctx", wantType: TypeContext,
 		},
@@ -344,6 +357,20 @@ import (
 )
 
 func SampleHandler(w http.ResponseWriter, req *http.Request) {}
+`,
+			wantName: "req", wantType: TypeHttpRequest,
+		},
+
+		{
+			name: "NamedHttp",
+			src: `
+package main
+
+import (
+	h "net/http"
+)
+
+func SampleHandler(w h.ResponseWriter, req *h.Request) {}
 `,
 			wantName: "req", wantType: TypeHttpRequest,
 		},
@@ -364,7 +391,7 @@ func SampleHandler(w http.ResponseWriter, req *http.Request) {}
 					break
 				}
 			}
-			name, gtype := parseParams(decl.Type)
+			name, gtype := parseParams(f.Imports, decl.Type)
 			if name != tt.wantName {
 				t.Errorf("parseParams() name = %q, want %q", name, tt.wantName)
 			}
@@ -374,3 +401,100 @@ func SampleHandler(w http.ResponseWriter, req *http.Request) {}
 		})
 	}
 }
+
+func Test_getImportName(t *testing.T) {
+	tests := []struct {
+		name string
+		src  string
+		typ  string
+		want string
+	}{
+		{
+			name: "Context",
+			src: `
+package main
+
+import (
+	"context"
+)
+
+func Hoge() {}`,
+			typ: TypeContext, want: "context",
+		},
+		{
+			name: "NoContext",
+			src: `
+package main
+
+import (
+	"net/http"
+)
+
+func Hoge() {}`,
+			typ: TypeContext, want: "context",
+		},
+		{
+			name: "NamedContext",
+			src: `
+package main
+
+import (
+	c "context"
+)
+
+func Hoge() {}`,
+			typ: TypeContext, want: "c",
+		},
+		{
+			name: "Http",
+			src: `
+package main
+
+import (
+	"net/http"
+)
+
+func Hoge() {}`,
+			typ: TypeHttpRequest, want: "http",
+		},
+		{
+			name: "NoHttp",
+			src: `
+package main
+
+import (
+	"context"
+)
+
+func Hoge() {}`,
+			typ: TypeHttpRequest, want: "http",
+		},
+		{
+			name: "NamedHttp",
+			src: `
+package main
+
+import (
+	myHttp "net/http"
+)
+
+func Hoge() {}`,
+			typ: TypeHttpRequest, want: "myHttp",
+		},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			fs := token.NewFileSet()
+			f, err := parser.ParseFile(fs, "sample.go", tt.src, parser.Mode(0))
+			if err != nil {
+				t.Fatal(err)
+			}
+			if got := getImportName(f.Imports, tt.typ); got != tt.want {
+				t.Errorf("getImportName() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}