github.com/keysonZZZ/kmg@v0.0.0-20151121023212-05317bfd7d39/kmgView/kmgGoTpl/addImport.go (about)

     1  package kmgGoTpl
     2  
     3  import (
     4  	"bytes"
     5  
     6  	"github.com/bronze1man/kmg/kmgStrings"
     7  )
     8  
     9  /*
    10  type importStatus int
    11  
    12  const (
    13  	importStatusNot          importStatus = 0
    14  	importStatusSpace1       importStatus = 1
    15  	importStatusImportToken  importStatus = 2
    16  	importStatusParentheses1 importStatus = 3
    17  	importStatusImportList   importStatus = 4
    18  	importStatusParentheses2 importStatus = 5
    19  )
    20  */
    21  
    22  type addImportStatus int
    23  
    24  const (
    25  	addImportStatusNot                addImportStatus = 0
    26  	addImportStatusPackageToken       addImportStatus = 1 //package
    27  	addImportStatusPackageLine        addImportStatus = 2
    28  	addImportStatusImportToken        addImportStatus = 3
    29  	addImportStatusImportParentheses1 addImportStatus = 4
    30  	addImportStatusImportParentheses2 addImportStatus = 6
    31  	addImportStatusImportImpossible   addImportStatus = 7
    32  )
    33  
    34  // 添加bytes的import项.
    35  // 这个是近似实现,golang的语法树确实有点复杂.
    36  func addImport(in []byte, pkgList []string) (out []byte) {
    37  	status := addImportStatus(0)
    38  	var hasFoundImport bool
    39  	var lastImportParenthesesPos int
    40  	outBuf := &bytes.Buffer{}
    41  	for pos := 0; pos < len(in); pos++ {
    42  		switch status {
    43  		case addImportStatusNot:
    44  			if bytes.HasPrefix(in[pos:], []byte("package")) {
    45  				status = addImportStatusPackageToken
    46  			}
    47  		case addImportStatusPackageToken:
    48  			if in[pos] == '\n' {
    49  				status = addImportStatusPackageLine
    50  			}
    51  		case addImportStatusPackageLine:
    52  			if bytes.HasPrefix(in[pos:], []byte("import")) {
    53  				status = addImportStatusImportToken
    54  			} else if bytes.HasPrefix(in[pos:], []byte("func ")) {
    55  				status = addImportStatusImportImpossible
    56  			} else if bytes.HasPrefix(in[pos:], []byte("type ")) {
    57  				status = addImportStatusImportImpossible
    58  			} else if bytes.HasPrefix(in[pos:], []byte("var ")) {
    59  				status = addImportStatusImportImpossible
    60  			}
    61  		case addImportStatusImportToken:
    62  			if in[pos] == '(' {
    63  				status = addImportStatusImportParentheses1
    64  				lastImportParenthesesPos = pos
    65  			}
    66  		case addImportStatusImportParentheses1:
    67  			if in[pos] == ')' {
    68  				status = addImportStatusImportParentheses2
    69  				hasFoundImport = true
    70  				var readedImportPathList []string
    71  				importPkgPath := bytes.Split(in[lastImportParenthesesPos+1:pos], []byte{'\n'})
    72  				for _, p := range importPkgPath {
    73  					p = bytes.TrimSpace(p)
    74  					if len(p) == 0 {
    75  						continue
    76  					}
    77  					readedImportPathList = append(readedImportPathList, string(p))
    78  				}
    79  				for _, pkg := range pkgList {
    80  					if !kmgStrings.IsInSlice(readedImportPathList, "\""+pkg+"\"") {
    81  						outBuf.WriteString("\"" + pkg + "\"\n")
    82  					}
    83  				}
    84  			}
    85  		}
    86  		outBuf.WriteByte(in[pos])
    87  	}
    88  	if !hasFoundImport || status == addImportStatusImportImpossible {
    89  		outBuf.Reset()
    90  		isLastPackageToken := false
    91  		hadAddImport := false
    92  		for pos := 0; pos < len(in); pos++ {
    93  			if !hadAddImport {
    94  				if bytes.HasPrefix(in[pos:], []byte("package ")) {
    95  					isLastPackageToken = true
    96  				} else if in[pos] == '\n' {
    97  					if isLastPackageToken {
    98  						outBuf.WriteString("\nimport (")
    99  						for _, pkg := range pkgList {
   100  							outBuf.WriteString("\"" + pkg + "\"\n")
   101  						}
   102  						outBuf.WriteString(")")
   103  						isLastPackageToken = false
   104  						hadAddImport = true
   105  					}
   106  				}
   107  			}
   108  			outBuf.WriteByte(in[pos])
   109  		}
   110  	}
   111  	return outBuf.Bytes()
   112  }