source.go 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170
  1. package source // import "gotest.tools/v3/internal/source"
  2. import (
  3. "bytes"
  4. "fmt"
  5. "go/ast"
  6. "go/format"
  7. "go/parser"
  8. "go/token"
  9. "os"
  10. "runtime"
  11. "strconv"
  12. "strings"
  13. "github.com/pkg/errors"
  14. )
  15. const baseStackIndex = 1
  16. // FormattedCallExprArg returns the argument from an ast.CallExpr at the
  17. // index in the call stack. The argument is formatted using FormatNode.
  18. func FormattedCallExprArg(stackIndex int, argPos int) (string, error) {
  19. args, err := CallExprArgs(stackIndex + 1)
  20. if err != nil {
  21. return "", err
  22. }
  23. if argPos >= len(args) {
  24. return "", errors.New("failed to find expression")
  25. }
  26. return FormatNode(args[argPos])
  27. }
  28. // CallExprArgs returns the ast.Expr slice for the args of an ast.CallExpr at
  29. // the index in the call stack.
  30. func CallExprArgs(stackIndex int) ([]ast.Expr, error) {
  31. _, filename, lineNum, ok := runtime.Caller(baseStackIndex + stackIndex)
  32. if !ok {
  33. return nil, errors.New("failed to get call stack")
  34. }
  35. debug("call stack position: %s:%d", filename, lineNum)
  36. node, err := getNodeAtLine(filename, lineNum)
  37. if err != nil {
  38. return nil, err
  39. }
  40. debug("found node: %s", debugFormatNode{node})
  41. return getCallExprArgs(node)
  42. }
  43. func getNodeAtLine(filename string, lineNum int) (ast.Node, error) {
  44. fileset := token.NewFileSet()
  45. astFile, err := parser.ParseFile(fileset, filename, nil, parser.AllErrors)
  46. if err != nil {
  47. return nil, errors.Wrapf(err, "failed to parse source file: %s", filename)
  48. }
  49. if node := scanToLine(fileset, astFile, lineNum); node != nil {
  50. return node, nil
  51. }
  52. if node := scanToDeferLine(fileset, astFile, lineNum); node != nil {
  53. node, err := guessDefer(node)
  54. if err != nil || node != nil {
  55. return node, err
  56. }
  57. }
  58. return nil, errors.Errorf(
  59. "failed to find an expression on line %d in %s", lineNum, filename)
  60. }
  61. func scanToLine(fileset *token.FileSet, node ast.Node, lineNum int) ast.Node {
  62. var matchedNode ast.Node
  63. ast.Inspect(node, func(node ast.Node) bool {
  64. switch {
  65. case node == nil || matchedNode != nil:
  66. return false
  67. case nodePosition(fileset, node).Line == lineNum:
  68. matchedNode = node
  69. return false
  70. }
  71. return true
  72. })
  73. return matchedNode
  74. }
  75. // In golang 1.9 the line number changed from being the line where the statement
  76. // ended to the line where the statement began.
  77. func nodePosition(fileset *token.FileSet, node ast.Node) token.Position {
  78. if goVersionBefore19 {
  79. return fileset.Position(node.End())
  80. }
  81. return fileset.Position(node.Pos())
  82. }
  83. // GoVersionLessThan returns true if runtime.Version() is semantically less than
  84. // version 1.minor.
  85. func GoVersionLessThan(minor int64) bool {
  86. version := runtime.Version()
  87. // not a release version
  88. if !strings.HasPrefix(version, "go") {
  89. return false
  90. }
  91. version = strings.TrimPrefix(version, "go")
  92. parts := strings.Split(version, ".")
  93. if len(parts) < 2 {
  94. return false
  95. }
  96. actual, err := strconv.ParseInt(parts[1], 10, 32)
  97. return err == nil && parts[0] == "1" && actual < minor
  98. }
  99. var goVersionBefore19 = GoVersionLessThan(9)
  100. func getCallExprArgs(node ast.Node) ([]ast.Expr, error) {
  101. visitor := &callExprVisitor{}
  102. ast.Walk(visitor, node)
  103. if visitor.expr == nil {
  104. return nil, errors.New("failed to find call expression")
  105. }
  106. debug("callExpr: %s", debugFormatNode{visitor.expr})
  107. return visitor.expr.Args, nil
  108. }
  109. type callExprVisitor struct {
  110. expr *ast.CallExpr
  111. }
  112. func (v *callExprVisitor) Visit(node ast.Node) ast.Visitor {
  113. if v.expr != nil || node == nil {
  114. return nil
  115. }
  116. debug("visit: %s", debugFormatNode{node})
  117. switch typed := node.(type) {
  118. case *ast.CallExpr:
  119. v.expr = typed
  120. return nil
  121. case *ast.DeferStmt:
  122. ast.Walk(v, typed.Call.Fun)
  123. return nil
  124. }
  125. return v
  126. }
  127. // FormatNode using go/format.Node and return the result as a string
  128. func FormatNode(node ast.Node) (string, error) {
  129. buf := new(bytes.Buffer)
  130. err := format.Node(buf, token.NewFileSet(), node)
  131. return buf.String(), err
  132. }
  133. var debugEnabled = os.Getenv("GOTESTTOOLS_DEBUG") != ""
  134. func debug(format string, args ...interface{}) {
  135. if debugEnabled {
  136. fmt.Fprintf(os.Stderr, "DEBUG: "+format+"\n", args...)
  137. }
  138. }
  139. type debugFormatNode struct {
  140. ast.Node
  141. }
  142. func (n debugFormatNode) String() string {
  143. out, err := FormatNode(n.Node)
  144. if err != nil {
  145. return fmt.Sprintf("failed to format %s: %s", n.Node, err)
  146. }
  147. return fmt.Sprintf("(%T) %s", n.Node, out)
  148. }