This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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