Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
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
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.