github.com/samlitowitz/goimportcycle@v1.0.9/internal/ast/dependency_visitor.go (about)

     1  package ast
     2  
     3  import (
     4  	"go/ast"
     5  	"go/token"
     6  	"path/filepath"
     7  	"strings"
     8  )
     9  
    10  type Package struct {
    11  	*ast.Package
    12  
    13  	DirName string
    14  }
    15  
    16  type File struct {
    17  	*ast.File
    18  
    19  	AbsPath string
    20  	DirName string
    21  }
    22  
    23  type ImportSpec struct {
    24  	*ast.ImportSpec
    25  
    26  	IsAliased bool
    27  	Alias     string
    28  }
    29  
    30  type FuncDecl struct {
    31  	*ast.FuncDecl
    32  
    33  	ReceiverName  string
    34  	QualifiedName string
    35  }
    36  
    37  type SelectorExpr struct {
    38  	*ast.SelectorExpr
    39  
    40  	ImportName string
    41  }
    42  
    43  func (decl FuncDecl) IsReceiver() bool {
    44  	return decl.Recv != nil
    45  }
    46  
    47  type DependencyVisitor struct {
    48  	out chan<- ast.Node
    49  
    50  	fileImports map[string]struct{}
    51  }
    52  
    53  func NewDependencyVisitor() (*DependencyVisitor, <-chan ast.Node) {
    54  	out := make(chan ast.Node)
    55  	v := &DependencyVisitor{
    56  		out: out,
    57  	}
    58  
    59  	return v, out
    60  }
    61  
    62  func (v *DependencyVisitor) Visit(node ast.Node) ast.Visitor {
    63  	switch node := node.(type) {
    64  	case *ast.Package:
    65  		v.emitPackageAndFiles(node)
    66  
    67  	case *ast.File:
    68  		v.fileImports = make(map[string]struct{})
    69  
    70  	case *ast.ImportSpec:
    71  		v.emitImportSpec(node)
    72  
    73  	case *ast.FuncDecl:
    74  		v.emitFuncDecl(node)
    75  
    76  	case *ast.GenDecl:
    77  		switch node.Tok {
    78  		case token.CONST:
    79  			fallthrough
    80  		case token.TYPE:
    81  			fallthrough
    82  		case token.VAR:
    83  			v.out <- node
    84  		}
    85  
    86  	case *ast.SelectorExpr:
    87  		// only references to external packages
    88  		if node.X == nil {
    89  			return v
    90  		}
    91  
    92  		impName := ""
    93  		switch x := node.X.(type) {
    94  		case *ast.Ident:
    95  			impName = x.String()
    96  		}
    97  
    98  		// if the "import name" is actually a variable and not a package, skip it
    99  		if _, ok := v.fileImports[impName]; !ok {
   100  			return v
   101  		}
   102  
   103  		v.out <- &SelectorExpr{
   104  			SelectorExpr: node,
   105  			ImportName:   impName,
   106  		}
   107  	}
   108  	return v
   109  }
   110  
   111  func (v *DependencyVisitor) emitPackageAndFiles(node *ast.Package) {
   112  	var setImportPathAndEmitPackage bool
   113  	var dirName string
   114  	for filename, astFile := range node.Files {
   115  		absPath, err := filepath.Abs(filename)
   116  		if err != nil {
   117  			continue
   118  		}
   119  		if !setImportPathAndEmitPackage {
   120  			dirName, _ = filepath.Split(absPath)
   121  			dirName = strings.TrimRight(dirName, "/")
   122  			v.out <- &Package{
   123  				Package: node,
   124  				DirName: dirName,
   125  			}
   126  			setImportPathAndEmitPackage = true
   127  		}
   128  
   129  		v.out <- &File{
   130  			File:    astFile,
   131  			AbsPath: absPath,
   132  			DirName: dirName,
   133  		}
   134  	}
   135  }
   136  
   137  func (v *DependencyVisitor) emitImportSpec(node *ast.ImportSpec) {
   138  	node.Path.Value = strings.Trim(node.Path.Value, "\"")
   139  	pieces := strings.Split(node.Path.Value, "/")
   140  	name := pieces[len(pieces)-1]
   141  
   142  	isAliased := node.Name != nil
   143  	alias := ""
   144  
   145  	if isAliased {
   146  		alias = node.Name.String()
   147  		node.Name.Name = name
   148  		v.fileImports[alias] = struct{}{}
   149  	}
   150  
   151  	if !isAliased {
   152  		node.Name = &ast.Ident{
   153  			Name: name,
   154  		}
   155  		v.fileImports[name] = struct{}{}
   156  	}
   157  
   158  	v.out <- &ImportSpec{
   159  		ImportSpec: node,
   160  		IsAliased:  isAliased,
   161  		Alias:      alias,
   162  	}
   163  }
   164  
   165  func (v *DependencyVisitor) emitFuncDecl(node *ast.FuncDecl) {
   166  	receiverName := ""
   167  	qualifiedName := node.Name.String()
   168  
   169  	if node.Recv != nil {
   170  		// TODO: don't emit receiver functions/methods? we don't need them
   171  		var typName string
   172  		switch expr := node.Recv.List[0].Type.(type) {
   173  		case *ast.Ident:
   174  			typName = expr.String()
   175  		case *ast.StarExpr:
   176  			if expr.X == nil {
   177  				// panic error, invalid receiver method
   178  			}
   179  			ident, ok := expr.X.(*ast.Ident)
   180  			if !ok {
   181  				// panic error, invalid receiver method
   182  			}
   183  			typName = ident.String()
   184  		default:
   185  			// panic error, invalid receiver method
   186  		}
   187  		receiverName = typName
   188  		qualifiedName = typName + "." + node.Name.String()
   189  	}
   190  
   191  	v.out <- &FuncDecl{
   192  		FuncDecl:      node,
   193  		ReceiverName:  receiverName,
   194  		QualifiedName: qualifiedName,
   195  	}
   196  }
   197  
   198  func (v *DependencyVisitor) Close() {
   199  	close(v.out)
   200  }