printer.go 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168
  1. package check
  2. import (
  3. "bytes"
  4. "go/ast"
  5. "go/parser"
  6. "go/printer"
  7. "go/token"
  8. "os"
  9. )
  10. func indent(s, with string) (r string) {
  11. eol := true
  12. for i := 0; i != len(s); i++ {
  13. c := s[i]
  14. switch {
  15. case eol && c == '\n' || c == '\r':
  16. case c == '\n' || c == '\r':
  17. eol = true
  18. case eol:
  19. eol = false
  20. s = s[:i] + with + s[i:]
  21. i += len(with)
  22. }
  23. }
  24. return s
  25. }
  26. func printLine(filename string, line int) (string, error) {
  27. fset := token.NewFileSet()
  28. file, err := os.Open(filename)
  29. if err != nil {
  30. return "", err
  31. }
  32. fnode, err := parser.ParseFile(fset, filename, file, parser.ParseComments)
  33. if err != nil {
  34. return "", err
  35. }
  36. config := &printer.Config{Mode: printer.UseSpaces, Tabwidth: 4}
  37. lp := &linePrinter{fset: fset, fnode: fnode, line: line, config: config}
  38. ast.Walk(lp, fnode)
  39. result := lp.output.Bytes()
  40. // Comments leave \n at the end.
  41. n := len(result)
  42. for n > 0 && result[n-1] == '\n' {
  43. n--
  44. }
  45. return string(result[:n]), nil
  46. }
  47. type linePrinter struct {
  48. config *printer.Config
  49. fset *token.FileSet
  50. fnode *ast.File
  51. line int
  52. output bytes.Buffer
  53. stmt ast.Stmt
  54. }
  55. func (lp *linePrinter) emit() bool {
  56. if lp.stmt != nil {
  57. lp.trim(lp.stmt)
  58. lp.printWithComments(lp.stmt)
  59. lp.stmt = nil
  60. return true
  61. }
  62. return false
  63. }
  64. func (lp *linePrinter) printWithComments(n ast.Node) {
  65. nfirst := lp.fset.Position(n.Pos()).Line
  66. nlast := lp.fset.Position(n.End()).Line
  67. for _, g := range lp.fnode.Comments {
  68. cfirst := lp.fset.Position(g.Pos()).Line
  69. clast := lp.fset.Position(g.End()).Line
  70. if clast == nfirst-1 && lp.fset.Position(n.Pos()).Column == lp.fset.Position(g.Pos()).Column {
  71. for _, c := range g.List {
  72. lp.output.WriteString(c.Text)
  73. lp.output.WriteByte('\n')
  74. }
  75. }
  76. if cfirst >= nfirst && cfirst <= nlast && n.End() <= g.List[0].Slash {
  77. // The printer will not include the comment if it starts past
  78. // the node itself. Trick it into printing by overlapping the
  79. // slash with the end of the statement.
  80. g.List[0].Slash = n.End() - 1
  81. }
  82. }
  83. node := &printer.CommentedNode{n, lp.fnode.Comments}
  84. lp.config.Fprint(&lp.output, lp.fset, node)
  85. }
  86. func (lp *linePrinter) Visit(n ast.Node) (w ast.Visitor) {
  87. if n == nil {
  88. if lp.output.Len() == 0 {
  89. lp.emit()
  90. }
  91. return nil
  92. }
  93. first := lp.fset.Position(n.Pos()).Line
  94. last := lp.fset.Position(n.End()).Line
  95. if first <= lp.line && last >= lp.line {
  96. // Print the innermost statement containing the line.
  97. if stmt, ok := n.(ast.Stmt); ok {
  98. if _, ok := n.(*ast.BlockStmt); !ok {
  99. lp.stmt = stmt
  100. }
  101. }
  102. if first == lp.line && lp.emit() {
  103. return nil
  104. }
  105. return lp
  106. }
  107. return nil
  108. }
  109. func (lp *linePrinter) trim(n ast.Node) bool {
  110. stmt, ok := n.(ast.Stmt)
  111. if !ok {
  112. return true
  113. }
  114. line := lp.fset.Position(n.Pos()).Line
  115. if line != lp.line {
  116. return false
  117. }
  118. switch stmt := stmt.(type) {
  119. case *ast.IfStmt:
  120. stmt.Body = lp.trimBlock(stmt.Body)
  121. case *ast.SwitchStmt:
  122. stmt.Body = lp.trimBlock(stmt.Body)
  123. case *ast.TypeSwitchStmt:
  124. stmt.Body = lp.trimBlock(stmt.Body)
  125. case *ast.CaseClause:
  126. stmt.Body = lp.trimList(stmt.Body)
  127. case *ast.CommClause:
  128. stmt.Body = lp.trimList(stmt.Body)
  129. case *ast.BlockStmt:
  130. stmt.List = lp.trimList(stmt.List)
  131. }
  132. return true
  133. }
  134. func (lp *linePrinter) trimBlock(stmt *ast.BlockStmt) *ast.BlockStmt {
  135. if !lp.trim(stmt) {
  136. return lp.emptyBlock(stmt)
  137. }
  138. stmt.Rbrace = stmt.Lbrace
  139. return stmt
  140. }
  141. func (lp *linePrinter) trimList(stmts []ast.Stmt) []ast.Stmt {
  142. for i := 0; i != len(stmts); i++ {
  143. if !lp.trim(stmts[i]) {
  144. stmts[i] = lp.emptyStmt(stmts[i])
  145. break
  146. }
  147. }
  148. return stmts
  149. }
  150. func (lp *linePrinter) emptyStmt(n ast.Node) *ast.ExprStmt {
  151. return &ast.ExprStmt{&ast.Ellipsis{n.Pos(), nil}}
  152. }
  153. func (lp *linePrinter) emptyBlock(n ast.Node) *ast.BlockStmt {
  154. p := n.Pos()
  155. return &ast.BlockStmt{p, []ast.Stmt{lp.emptyStmt(n)}, p}
  156. }