Skip to content

Commit

Permalink
Add import aliases
Browse files Browse the repository at this point in the history
  • Loading branch information
@DimitarPetrov
DimitarPetrov committed Sep 6, 2020
1 parent b60eec9 commit 1b71478
Show file tree
Hide file tree
Showing 10 changed files with 70 additions and 64 deletions.
2 changes: 1 addition & 1 deletion cmd/apply.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,6 @@ func (ac *ApplyCmd) Run() error {
if err != nil {
return err
}
return ac.importsGroomer.RemoveUnusedImportFromDirectory(path, []string{"fmt", "runtime", "rand"})
return ac.importsGroomer.RemoveUnusedImportFromDirectory(path, map[string]string{"fmt":"", "runtime":"rt", "rand":""})
})
}
2 changes: 1 addition & 1 deletion cmd/revert.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,6 @@ func (rc *RevertCmd) Run() error {
if err != nil {
return err
}
return rc.importsGroomer.RemoveUnusedImportFromDirectory(path, []string{"fmt", "runtime", "rand"})
return rc.importsGroomer.RemoveUnusedImportFromDirectory(path, map[string]string{"fmt":"", "runtime":"rt", "rand":""})
})
}
4 changes: 3 additions & 1 deletion tracing/deinstrument.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,9 @@ func (cd *codeDeinstrumenter) DeinstrumentFile(fset *token.FileSet, file *ast.Fi

if checkInstrumentationStatementsIntegrity(t) {
t.Body.List = t.Body.List[instrumentationStmtsCount:]
t.Body.List[0].Decorations().Before = dst.None
if len(t.Body.List) > 0 {
t.Body.List[0].Decorations().Before = dst.None
}
}
}
}
Expand Down
4 changes: 2 additions & 2 deletions tracing/deinstrument_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ func TestDeinstrumentFile(t *testing.T) {
}

var buff2 bytes.Buffer
if err := NewImportsGroomer().RemoveUnusedImportFromFile(fset, file, &buff2, []string{"fmt", "runtime", "rand"}); err != nil {
if err := NewImportsGroomer().RemoveUnusedImportFromFile(fset, file, &buff2, map[string]string{"fmt":"", "runtime":"rt", "rand":""}); err != nil {
t.Fatal(err)
}

Expand Down Expand Up @@ -87,7 +87,7 @@ func TestDeinstrumentDirectory(t *testing.T) {
t.Fatal(err)
}

if err := NewImportsGroomer().RemoveUnusedImportFromDirectory("test", []string{"fmt", "runtime", "rand"}); err != nil {
if err := NewImportsGroomer().RemoveUnusedImportFromDirectory("test", map[string]string{"fmt":"", "runtime":"rt", "rand":""}); err != nil {
t.Fatal(err)
}

Expand Down
2 changes: 1 addition & 1 deletion tracing/instrument.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,7 @@ func (ci *codeInstrumenter) InstrumentPackage(fset *token.FileSet, pkg *ast.Pack

func (ci *codeInstrumenter) InstrumentFile(fset *token.FileSet, file *ast.File, out io.Writer) error {
astutil.AddImport(fset, file, "fmt")
astutil.AddImport(fset, file, "runtime")
astutil.AddNamedImport(fset, file, "rt", "runtime")
astutil.AddImport(fset, file, "rand")

// Needed because ast does not support floating comments and deletes them.
Expand Down
92 changes: 46 additions & 46 deletions tracing/instrument_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,19 +29,19 @@ const resultCodeWithoutImports = `package a
import (
"fmt"
"rand"
"runtime"
rt "runtime"
)
func test(i int, b bool) int {
/* prinTracer */
funcName := "test"
caller := "unknown"
if funcPC, _, _, ok := runtime.Caller(0); ok {
funcName = runtime.FuncForPC(funcPC).Name()
if funcPC, _, _, ok := rt.Caller(0); ok {
funcName = rt.FuncForPC(funcPC).Name()
}
if callerPC, _, _, ok := runtime.Caller(1); ok {
caller = runtime.FuncForPC(callerPC).Name()
if callerPC, _, _, ok := rt.Caller(1); ok {
caller = rt.FuncForPC(callerPC).Name()
}
idBytes := make([]byte, 16)
_, _ = rand.Read(idBytes)
Expand All @@ -60,11 +60,11 @@ func main() {
/* prinTracer */
funcName := "main"
caller := "unknown"
if funcPC, _, _, ok := runtime.Caller(0); ok {
funcName = runtime.FuncForPC(funcPC).Name()
if funcPC, _, _, ok := rt.Caller(0); ok {
funcName = rt.FuncForPC(funcPC).Name()
}
if callerPC, _, _, ok := runtime.Caller(1); ok {
caller = runtime.FuncForPC(callerPC).Name()
if callerPC, _, _, ok := rt.Caller(1); ok {
caller = rt.FuncForPC(callerPC).Name()
}
idBytes := make([]byte, 16)
_, _ = rand.Read(idBytes)
Expand All @@ -81,19 +81,19 @@ const editedResultCodeWithoutImports = `package a
import (
"fmt"
"rand"
"runtime"
rt "runtime"
)
func test(i int, b bool) int {
/* prinTracer */
funcName := "test2"
caller := "unknown2"
if funcPC, _, _, ok := runtime.Caller(0); ok {
funcName = runtime.FuncForPC(funcPC).Name()
if funcPC, _, _, ok := rt.Caller(0); ok {
funcName = rt.FuncForPC(funcPC).Name()
}
if callerPC, _, _, ok := runtime.Caller(1); ok {
caller = runtime.FuncForPC(callerPC).Name()
if callerPC, _, _, ok := rt.Caller(1); ok {
caller = rt.FuncForPC(callerPC).Name()
}
fmt.Println("test")
idBytes := make([]byte, 16)
Expand All @@ -112,11 +112,11 @@ func main() {
funcName := "main"
caller := "unknown"
if funcPC, _, _, ok := runtime.Caller(0); ok {
funcName = runtime.FuncForPC(funcPC).Name()
if funcPC, _, _, ok := rt.Caller(0); ok {
funcName = rt.FuncForPC(funcPC).Name()
}
if callerPC, _, _, ok := runtime.Caller(1); ok {
caller = runtime.FuncForPC(callerPC).Name()
if callerPC, _, _, ok := rt.Caller(1); ok {
caller = rt.FuncForPC(callerPC).Name()
}
idBytes := make([]byte, 16)
_, _ = rand.Read(idBytes)
Expand Down Expand Up @@ -151,19 +151,19 @@ const resultCodeWithFmtImport = `package a
import (
"fmt"
"rand"
"runtime"
rt "runtime"
)
func test(i int, b bool) int {
/* prinTracer */
funcName := "test"
caller := "unknown"
if funcPC, _, _, ok := runtime.Caller(0); ok {
funcName = runtime.FuncForPC(funcPC).Name()
if funcPC, _, _, ok := rt.Caller(0); ok {
funcName = rt.FuncForPC(funcPC).Name()
}
if callerPC, _, _, ok := runtime.Caller(1); ok {
caller = runtime.FuncForPC(callerPC).Name()
if callerPC, _, _, ok := rt.Caller(1); ok {
caller = rt.FuncForPC(callerPC).Name()
}
idBytes := make([]byte, 16)
_, _ = rand.Read(idBytes)
Expand All @@ -182,11 +182,11 @@ func main() {
/* prinTracer */
funcName := "main"
caller := "unknown"
if funcPC, _, _, ok := runtime.Caller(0); ok {
funcName = runtime.FuncForPC(funcPC).Name()
if funcPC, _, _, ok := rt.Caller(0); ok {
funcName = rt.FuncForPC(funcPC).Name()
}
if callerPC, _, _, ok := runtime.Caller(1); ok {
caller = runtime.FuncForPC(callerPC).Name()
if callerPC, _, _, ok := rt.Caller(1); ok {
caller = rt.FuncForPC(callerPC).Name()
}
idBytes := make([]byte, 16)
_, _ = rand.Read(idBytes)
Expand Down Expand Up @@ -224,7 +224,7 @@ const resultCodeWithMultipleImports = `package a
import (
"fmt"
"rand"
"runtime"
rt "runtime"
"strconv"
)
Expand All @@ -233,11 +233,11 @@ func test(i int, b bool) int {
/* prinTracer */
funcName := "test"
caller := "unknown"
if funcPC, _, _, ok := runtime.Caller(0); ok {
funcName = runtime.FuncForPC(funcPC).Name()
if funcPC, _, _, ok := rt.Caller(0); ok {
funcName = rt.FuncForPC(funcPC).Name()
}
if callerPC, _, _, ok := runtime.Caller(1); ok {
caller = runtime.FuncForPC(callerPC).Name()
if callerPC, _, _, ok := rt.Caller(1); ok {
caller = rt.FuncForPC(callerPC).Name()
}
idBytes := make([]byte, 16)
_, _ = rand.Read(idBytes)
Expand All @@ -256,11 +256,11 @@ func main() {
/* prinTracer */
funcName := "main"
caller := "unknown"
if funcPC, _, _, ok := runtime.Caller(0); ok {
funcName = runtime.FuncForPC(funcPC).Name()
if funcPC, _, _, ok := rt.Caller(0); ok {
funcName = rt.FuncForPC(funcPC).Name()
}
if callerPC, _, _, ok := runtime.Caller(1); ok {
caller = runtime.FuncForPC(callerPC).Name()
if callerPC, _, _, ok := rt.Caller(1); ok {
caller = rt.FuncForPC(callerPC).Name()
}
idBytes := make([]byte, 16)
_, _ = rand.Read(idBytes)
Expand Down Expand Up @@ -297,7 +297,7 @@ const resultCodeWithImportsWithoutFmt = `package a
import (
"fmt"
"rand"
"runtime"
rt "runtime"
"strconv"
)
Expand All @@ -306,11 +306,11 @@ func test(i int, b bool) int {
/* prinTracer */
funcName := "test"
caller := "unknown"
if funcPC, _, _, ok := runtime.Caller(0); ok {
funcName = runtime.FuncForPC(funcPC).Name()
if funcPC, _, _, ok := rt.Caller(0); ok {
funcName = rt.FuncForPC(funcPC).Name()
}
if callerPC, _, _, ok := runtime.Caller(1); ok {
caller = runtime.FuncForPC(callerPC).Name()
if callerPC, _, _, ok := rt.Caller(1); ok {
caller = rt.FuncForPC(callerPC).Name()
}
idBytes := make([]byte, 16)
_, _ = rand.Read(idBytes)
Expand All @@ -329,11 +329,11 @@ func main() {
/* prinTracer */
funcName := "main"
caller := "unknown"
if funcPC, _, _, ok := runtime.Caller(0); ok {
funcName = runtime.FuncForPC(funcPC).Name()
if funcPC, _, _, ok := rt.Caller(0); ok {
funcName = rt.FuncForPC(funcPC).Name()
}
if callerPC, _, _, ok := runtime.Caller(1); ok {
caller = runtime.FuncForPC(callerPC).Name()
if callerPC, _, _, ok := rt.Caller(1); ok {
caller = rt.FuncForPC(callerPC).Name()
}
idBytes := make([]byte, 16)
_, _ = rand.Read(idBytes)
Expand All @@ -358,7 +358,7 @@ const resultCodeWithoutFunction = `package a
import (
"fmt"
"rand"
"runtime"
rt "runtime"
)
type test struct {
Expand Down
6 changes: 3 additions & 3 deletions tracing/interfaces.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ type CodeDeinstrumenter interface {

//go:generate counterfeiter . ImportsGroomer
type ImportsGroomer interface {
RemoveUnusedImportFromFile(fset *token.FileSet, file *ast.File, out io.Writer, importsToRemove []string) error
RemoveUnusedImportFromPackage(fset *token.FileSet, pkg *ast.Package, importsToRemove []string) error
RemoveUnusedImportFromDirectory(path string, importsToRemove []string) error
RemoveUnusedImportFromFile(fset *token.FileSet, file *ast.File, out io.Writer, importsToRemove map[string]string) error
RemoveUnusedImportFromPackage(fset *token.FileSet, pkg *ast.Package, importsToRemove map[string]string) error
RemoveUnusedImportFromDirectory(path string, importsToRemove map[string]string) error
}
10 changes: 5 additions & 5 deletions tracing/unused_imports.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ func NewImportsGroomer() ImportsGroomer {
return &importsGroomer{}
}

func (ig *importsGroomer) RemoveUnusedImportFromDirectory(path string, importsToRemove []string) error {
func (ig *importsGroomer) RemoveUnusedImportFromDirectory(path string, importsToRemove map[string]string) error {
fset := token.NewFileSet()
filter := func(info os.FileInfo) bool {
return testsFilter(info) && generatedFilter(path, info)
Expand All @@ -36,7 +36,7 @@ func (ig *importsGroomer) RemoveUnusedImportFromDirectory(path string, importsTo
return nil
}

func (ig *importsGroomer) RemoveUnusedImportFromPackage(fset *token.FileSet, pkg *ast.Package, importsToRemove []string) error {
func (ig *importsGroomer) RemoveUnusedImportFromPackage(fset *token.FileSet, pkg *ast.Package, importsToRemove map[string]string) error {
for fileName, file := range pkg.Files {
sourceFile, err := os.OpenFile(fileName, os.O_TRUNC|os.O_WRONLY, 0664)
if err != nil {
Expand All @@ -49,10 +49,10 @@ func (ig *importsGroomer) RemoveUnusedImportFromPackage(fset *token.FileSet, pkg
return nil
}

func (ig *importsGroomer) RemoveUnusedImportFromFile(fset *token.FileSet, file *ast.File, out io.Writer, importsToRemove []string) error {
for _, importToRemove := range importsToRemove {
func (ig *importsGroomer) RemoveUnusedImportFromFile(fset *token.FileSet, file *ast.File, out io.Writer, importsToRemove map[string]string) error {
for importToRemove, alias := range importsToRemove {
if !astutil.UsesImport(file, importToRemove) {
astutil.DeleteImport(fset, file, importToRemove)
astutil.DeleteNamedImport(fset, file, alias, importToRemove)
}
}
// Needed because ast does not support floating comments and deletes them.
Expand Down
4 changes: 2 additions & 2 deletions tracing/unused_imports_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ func TestRemoveUnusedImportsFromFile(t *testing.T) {
t.Fatal(err)
}
var buff bytes.Buffer
if err := NewImportsGroomer().RemoveUnusedImportFromFile(fset, file, &buff, []string{"fmt", "runtime", "rand"}); err != nil {
if err := NewImportsGroomer().RemoveUnusedImportFromFile(fset, file, &buff, map[string]string{"fmt":"", "runtime":"rt", "rand":""}); err != nil {
t.Fatal(err)
}

Expand Down Expand Up @@ -63,7 +63,7 @@ func TestRemoveUnusedImportsFromDirectory(t *testing.T) {
i++
}

if err := NewImportsGroomer().RemoveUnusedImportFromDirectory("test", []string{"fmt", "runtime", "rand"}); err != nil {
if err := NewImportsGroomer().RemoveUnusedImportFromDirectory("test", map[string]string{"fmt":"", "runtime":"rt", "rand":""}); err != nil {
t.Fatal(err)
}

Expand Down
8 changes: 6 additions & 2 deletions tracing/util.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -165,7 +165,7 @@ func newGetFuncNameIfStatement(funcIndex, funcPcVarName, funcNameVarName string)
&dst.CallExpr{
Fun: &dst.SelectorExpr{
X: &dst.Ident{
Name: "runtime",
Name: "rt",
},
Sel: &dst.Ident{
Name: "Caller",
Expand Down Expand Up @@ -198,7 +198,7 @@ func newGetFuncNameIfStatement(funcIndex, funcPcVarName, funcNameVarName string)
X: &dst.CallExpr{
Fun: &dst.SelectorExpr{
X: &dst.Ident{
Name: "runtime",
Name: "rt",
},
Sel: &dst.Ident{
Name: "FuncForPC",
Expand Down Expand Up @@ -253,6 +253,8 @@ func newMakeByteSliceStmt() *dst.AssignStmt {
}
}

// Returns dst statement like:
// _, _ = rand.Read(idBytes)
func newRandReadStmt() *dst.AssignStmt {
return &dst.AssignStmt{
Lhs: []dst.Expr{
Expand Down Expand Up @@ -284,6 +286,8 @@ func newRandReadStmt() *dst.AssignStmt {
}
}

// Returns dst statement like:
// callID := fmt.Sprintf("%x-%x-%x-%x-%x", idBytes[0:4], idBytes[4:6], idBytes[6:8], idBytes[8:10], idBytes[10:])
func newParseUUIDFromByteSliceStmt(callIDVarName string) *dst.AssignStmt {
return &dst.AssignStmt{
Lhs: []dst.Expr{
Expand Down

0 comments on commit 1b71478

Please sign in to comment.