ast_expression_parser.go

package expression

import (
	"fmt"
	"go/ast"
	"go/parser"
)

// FindSelectors parses an expression and returns any selectors used in the expression
func FindSelectors(expr string) ([]string, error) {
	tree, err := parser.ParseExpr(expr)
	if err != nil {
		return nil, err
	}

	// check for single *ast.ExprStmt.X = *ast.Ident
	if ident, ok := tree.(*ast.Ident); ok {
		return []string{ident.Name}, nil
	}

	var v expressionVisitor
	ast.Walk(&v, tree)
	return v.selectors, nil
}

// expressionVisitor implements ast.Visitor and tracks state as it traverses the
// expression abstract syntax tree to extract selectors (ie 'pack.name')
type expressionVisitor struct {
	selectors []string

	// collect and endSelector is used for collecting parts of complex selectors
	// (ie 'a.b.c.d.e')
	collect     string
	endSelector string

	prev *ast.Ident
}

// Visit is called by ast.Walk on the nodes of the expression tree.
// An expression looks like this:
//
//     ast.ExprStmt from `parser.ParseExpr("my.deep.selector - 3")`
//      └── X: *ast.BinaryExpr (Op: `-`)
//          ├── X: *ast.SelectorExpr
//          |   ├── X: *ast.SelectorExpr
//          |   |   ├── X: *ast.Ident (Name: 'my')
//          |   |   └── Sel: *ast.Ident (Name: 'deep')
//          |   └── Sel: *ast.Ident (Name: 'selector')
//          └── Y: *ast.BasicLit (Kind: INT)
//
// Nodes are visited depth-first by ast.Walk
func (v *expressionVisitor) Visit(n ast.Node) ast.Visitor {
	if n == nil {
		return v
	}

	switch node := n.(type) {
	case *ast.SelectorExpr:
		switch x := node.X.(type) {
		case *ast.SelectorExpr:
			if v.collect != "" {
				v.collect = fmt.Sprintf("%s.%s", x.Sel.Name, v.collect)
			} else {
				v.collect, v.endSelector = x.Sel.Name, node.Sel.Name
			}

		case *ast.Ident:
			if v.collect != "" {
				// if a collection was started, assemble selector and reset
				v.selectors = append(v.selectors, fmt.Sprintf("%s.%s.%s",
					x.Name, v.collect, v.endSelector))
				v.collect, v.endSelector = "", ""
			} else {
				v.selectors = append(v.selectors,
					fmt.Sprintf("%s.%s", x.Name, node.Sel.Name))
			}
		}

	// handle standalone variables
	case *ast.BinaryExpr:
		if x, ok := node.X.(*ast.Ident); ok {
			v.selectors = append(v.selectors, x.Name)
		}
		if y, ok := node.Y.(*ast.Ident); ok {
			v.selectors = append(v.selectors, y.Name)
		}
	}

	return v
}