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