-
Notifications
You must be signed in to change notification settings - Fork 152
/
option_editor.go
130 lines (108 loc) · 3.5 KB
/
option_editor.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
package edit
import (
"fmt"
"github.com/influxdata/flux/ast"
)
// `OptionFn` is a function that, provided with an `OptionStatement`, returns
// an `Expression` or an error. It is used by `Option` functions to edit
// AST's options statements.
type OptionFn func(opt *ast.OptionStatement) (ast.Expression, error)
// `Option` passes the `OptionStatement` in the AST rooted at `node` that has the
// specified identifier to `fn`.
// The function can have side effects on the option statement
// and/or return a non-nil `Expression` that is set as value for the option.
// If the value returned by the edit function is `nil` (or an error is returned) no new value is set
// for the option statement (but any, maybe partial, side effect is applied).
// `Option` returns whether it could find and edit the option (possibly with errors) or not.
func Option(node ast.Node, optionIdentifier string, fn OptionFn) (bool, error) {
oe := &optionEditor{identifier: optionIdentifier, optionFn: fn, err: nil}
ast.Walk(oe, node)
if oe.err != nil {
return oe.found, oe.err
}
return oe.found, nil
}
// Creates an `OptionFn` for setting the value of an `OptionStatement`.
func OptionValueFn(expr ast.Expression) OptionFn {
return func(opt *ast.OptionStatement) (ast.Expression, error) {
return expr, nil
}
}
// Creates an `OptionFn` for updating the values of an `OptionStatement` that has an
// `ObjectExpression` as value. Returns error if the child of the option statement is not
// an object expression. If some key is not a property of the object it is added.
func OptionObjectFn(keyMap map[string]ast.Expression) OptionFn {
return func(opt *ast.OptionStatement) (ast.Expression, error) {
a, ok := opt.Assignment.(*ast.VariableAssignment)
if !ok {
return nil, fmt.Errorf("option assignment must be variable assignment")
}
obj, ok := a.Init.(*ast.ObjectExpression)
if !ok {
return nil, fmt.Errorf("value is is %s, not an object expression", a.Init.Type())
}
// check that every specified property exists in the object
found := make(map[string]bool, len(obj.Properties))
for _, p := range obj.Properties {
found[p.Key.Key()] = true
}
for k := range keyMap {
if !found[k] {
obj.Properties = append(obj.Properties, &ast.Property{
Key: &ast.Identifier{Name: k},
Value: keyMap[k],
})
}
}
for _, p := range obj.Properties {
exp, found := keyMap[p.Key.Key()]
if found {
p.Value = exp
}
}
return nil, nil
}
}
//Finds the `OptionStatement` with the specified `identifier` and updates its value.
//There shouldn't be more then one option statement with the same identifier
//in a valid query.
type optionEditor struct {
identifier string
optionFn OptionFn
err error
found bool
}
func (v *optionEditor) Visit(node ast.Node) ast.Visitor {
if os, ok := node.(*ast.OptionStatement); ok {
switch a := os.Assignment.(type) {
case *ast.VariableAssignment:
if a.ID.Name == v.identifier {
v.found = true
newInit, err := v.optionFn(os)
if err != nil {
v.err = err
} else if newInit != nil {
a.Init = newInit
}
return nil
}
case *ast.MemberAssignment:
id, ok := a.Member.Object.(*ast.Identifier)
if ok {
name := id.Name + "." + a.Member.Property.Key()
if name == v.identifier {
v.found = true
newInit, err := v.optionFn(os)
if err != nil {
v.err = err
} else if newInit != nil {
a.Init = newInit
}
return nil
}
}
}
}
return v
}
func (v *optionEditor) Done(node ast.Node) {}