From b2f1ca95272ebdfad496d71daf3167efc931caec Mon Sep 17 00:00:00 2001
From: Pablo Baeyens <pablo.baeyens@datadoghq.com>
Date: Fri, 8 Oct 2021 11:14:17 +0200
Subject: [PATCH 1/2] Skip generated files

---
 import.go                     | 30 +++++++++++++++++++++++++++++-
 import_test.go                | 10 ++++++++++
 testdata/codegen/generated.go | 16 ++++++++++++++++
 3 files changed, 55 insertions(+), 1 deletion(-)
 create mode 100644 testdata/codegen/generated.go

diff --git a/import.go b/import.go
index e63df28..eb026aa 100644
--- a/import.go
+++ b/import.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
+	"go/ast"
 	"go/parser"
 	"go/token"
 	"io/ioutil"
@@ -14,12 +15,33 @@ import (
 
 var (
 	errMainPackage = errors.New("failed to add import to a main package")
+	errGenerated   = errors.New("failed to add import to a generated file")
+	// Matches https://golang.org/s/generatedcode and cgo generated comment.
+	// Taken from https://github.com/golang/tools/blob/c5188f24a/refactor/rename/spec.go#L574-L576
+	generatedRx = regexp.MustCompile(`// .*DO NOT EDIT\.?`)
 )
 
+// generated reports whether ast.File is a generated file.
+// Taken from https://github.com/golang/tools/blob/c5188f24a/refactor/rename/spec.go#L578-L593
+func generated(pf *ast.File, tokenFile *token.File) bool {
+	// Iterate over the comments in the file
+	for _, commentGroup := range pf.Comments {
+		for _, comment := range commentGroup.List {
+			if matched := generatedRx.MatchString(comment.Text); matched {
+				// Check if comment is at the beginning of the line in source
+				if pos := tokenFile.Position(comment.Slash); pos.Column == 1 {
+					return true
+				}
+			}
+		}
+	}
+	return false
+}
+
 // addImportPath adds the vanity import path to a given go file.
 func addImportPath(absFilepath string, module string) (bool, []byte, error) {
 	fset := token.NewFileSet()
-	pf, err := parser.ParseFile(fset, absFilepath, nil, 0)
+	pf, err := parser.ParseFile(fset, absFilepath, nil, parser.ParseComments)
 	if err != nil {
 		return false, nil, fmt.Errorf("failed to parse the file %q: %v", absFilepath, err)
 	}
@@ -28,6 +50,12 @@ func addImportPath(absFilepath string, module string) (bool, []byte, error) {
 		return false, nil, errMainPackage
 	}
 
+	// Skip generated files.
+	tokenFile := fset.File(pf.Pos())
+	if generated(pf, tokenFile) {
+		return false, nil, errGenerated
+	}
+
 	content, err := ioutil.ReadFile(absFilepath)
 	if err != nil {
 		return false, nil, fmt.Errorf("failed to parse the file %q: %v", absFilepath, err)
diff --git a/import_test.go b/import_test.go
index a608f90..f03010d 100644
--- a/import_test.go
+++ b/import_test.go
@@ -20,6 +20,16 @@ func TestAddImportPathAddsVanityImport(t *testing.T) {
 	assert.Equal(t, "package leftpad // import \"mypackage\"", string(newContent[15:52]))
 }
 
+func TestAddImportAutogenerated(t *testing.T) {
+	cwd, _ := os.Getwd()
+	hasChanged, _, err := addImportPath(
+		cwd+"/testdata/codegen/generated.go",
+		"codegen")
+
+	assert.Equal(t, errGenerated, err)
+	assert.False(t, hasChanged)
+}
+
 func TestAddImportPathFixesTheVanityImport(t *testing.T) {
 	cwd, _ := os.Getwd()
 	hasChanged, newContent, err := addImportPath(
diff --git a/testdata/codegen/generated.go b/testdata/codegen/generated.go
new file mode 100644
index 0000000..b9cf200
--- /dev/null
+++ b/testdata/codegen/generated.go
@@ -0,0 +1,16 @@
+// Code generated by Thrift Compiler (0.14.1). DO NOT EDIT.
+
+package codegen
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"time"
+)
+
+// (needed to ensure safety because of naive import list construction.)
+var _ = fmt.Printf
+var _ = context.Background
+var _ = time.Now
+var _ = bytes.Equal

From 22b934650c07426380c2779670df31d407c1c70c Mon Sep 17 00:00:00 2001
From: Pablo Baeyens <pablo.baeyens@datadoghq.com>
Date: Fri, 8 Oct 2021 16:59:55 +0200
Subject: [PATCH 2/2] Address review comment

---
 import.go | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/import.go b/import.go
index eb026aa..bd4dbd2 100644
--- a/import.go
+++ b/import.go
@@ -21,9 +21,9 @@ var (
 	generatedRx = regexp.MustCompile(`// .*DO NOT EDIT\.?`)
 )
 
-// generated reports whether ast.File is a generated file.
+// isGeneratedFile reports whether ast.File is a generated file.
 // Taken from https://github.com/golang/tools/blob/c5188f24a/refactor/rename/spec.go#L578-L593
-func generated(pf *ast.File, tokenFile *token.File) bool {
+func isGeneratedFile(pf *ast.File, tokenFile *token.File) bool {
 	// Iterate over the comments in the file
 	for _, commentGroup := range pf.Comments {
 		for _, comment := range commentGroup.List {
@@ -52,7 +52,7 @@ func addImportPath(absFilepath string, module string) (bool, []byte, error) {
 
 	// Skip generated files.
 	tokenFile := fset.File(pf.Pos())
-	if generated(pf, tokenFile) {
+	if isGeneratedFile(pf, tokenFile) {
 		return false, nil, errGenerated
 	}