package starlark import ( "fmt" "sort" "time" "go.starlark.net/starlark" "github.com/influxdata/telegraf" "github.com/influxdata/telegraf/metric" ) func newMetric(_ *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { var ( name starlark.String tags, fields starlark.Value ) if err := starlark.UnpackArgs("Metric", args, kwargs, "name", &name, "tags?", &tags, "fields?", &fields); err != nil { return nil, err } allFields, err := toFields(fields) if err != nil { return nil, err } allTags, err := toTags(tags) if err != nil { return nil, err } m := metric.New(string(name), allTags, allFields, time.Now()) return &Metric{metric: m}, nil } func toString(value starlark.Value, errorMsg string) (string, error) { if value, ok := value.(starlark.String); ok { return string(value), nil } return "", fmt.Errorf(errorMsg, value) } func items(value starlark.Value, errorMsg string) ([]starlark.Tuple, error) { if iter, ok := value.(starlark.IterableMapping); ok { return iter.Items(), nil } return nil, fmt.Errorf(errorMsg, value) } func deepcopy(_ *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { var sm *Metric var track bool if err := starlark.UnpackArgs("deepcopy", args, kwargs, "source", &sm, "track?", &track); err != nil { return nil, err } // In case we copy a tracking metric but do not want to track the result, // we have to strip the tracking information. This can be done by unwrapping // the metric. if tm, ok := sm.metric.(telegraf.TrackingMetric); ok && !track { return &Metric{metric: tm.Unwrap().Copy()}, nil } // Copy the whole metric including potential tracking information return &Metric{metric: sm.metric.Copy()}, nil } // catch(f) evaluates f() and returns its evaluation error message // if it failed or None if it succeeded. func catch(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { var fn starlark.Callable if err := starlark.UnpackArgs("catch", args, kwargs, "fn", &fn); err != nil { return nil, err } if _, err := starlark.Call(thread, fn, nil, nil); err != nil { //nolint:nilerr // nil returned on purpose, error put inside starlark.Value return starlark.String(err.Error()), nil } return starlark.None, nil } type builtinMethod func(b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) func builtinAttr(recv starlark.Value, name string, methods map[string]builtinMethod) (starlark.Value, error) { method := methods[name] if method == nil { return starlark.None, fmt.Errorf("no such method %q", name) } // Allocate a closure over 'method'. impl := func(_ *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { return method(b, args, kwargs) } return starlark.NewBuiltin(name, impl).BindReceiver(recv), nil } func builtinAttrNames(methods map[string]builtinMethod) []string { names := make([]string, 0, len(methods)) for name := range methods { names = append(names, name) } sort.Strings(names) return names } // --- dictionary methods --- // https://github.com/google/starlark-go/blob/master/doc/spec.md#dict·clear func dictClear(b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { if err := starlark.UnpackPositionalArgs(b.Name(), args, kwargs, 0); err != nil { return starlark.None, fmt.Errorf("%s: %w", b.Name(), err) } type HasClear interface { Clear() error } return starlark.None, b.Receiver().(HasClear).Clear() } // https://github.com/google/starlark-go/blob/master/doc/spec.md#dict·pop func dictPop(b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { var k, d starlark.Value if err := starlark.UnpackPositionalArgs(b.Name(), args, kwargs, 1, &k, &d); err != nil { return starlark.None, fmt.Errorf("%s: %w", b.Name(), err) } type HasDelete interface { Delete(k starlark.Value) (starlark.Value, bool, error) } if v, found, err := b.Receiver().(HasDelete).Delete(k); err != nil { return starlark.None, fmt.Errorf("%s: %w", b.Name(), err) // dict is frozen or key is unhashable } else if found { return v, nil } else if d != nil { return d, nil } return starlark.None, fmt.Errorf("%s: missing key", b.Name()) } // https://github.com/google/starlark-go/blob/master/doc/spec.md#dict·popitem func dictPopitem(b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { if err := starlark.UnpackPositionalArgs(b.Name(), args, kwargs, 0); err != nil { return starlark.None, fmt.Errorf("%s: %w", b.Name(), err) } type HasPopItem interface { PopItem() (starlark.Value, error) } return b.Receiver().(HasPopItem).PopItem() } // https://github.com/google/starlark-go/blob/master/doc/spec.md#dict·get func dictGet(b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { var key, dflt starlark.Value if err := starlark.UnpackPositionalArgs(b.Name(), args, kwargs, 1, &key, &dflt); err != nil { return starlark.None, fmt.Errorf("%s: %w", b.Name(), err) } if v, ok, err := b.Receiver().(starlark.Mapping).Get(key); err != nil { return starlark.None, fmt.Errorf("%s: %w", b.Name(), err) } else if ok { return v, nil } else if dflt != nil { return dflt, nil } return starlark.None, nil } // https://github.com/google/starlark-go/blob/master/doc/spec.md#dict·setdefault func dictSetdefault(b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { var key, dflt starlark.Value = nil, starlark.None if err := starlark.UnpackPositionalArgs(b.Name(), args, kwargs, 1, &key, &dflt); err != nil { return starlark.None, fmt.Errorf("%s: %w", b.Name(), err) } recv := b.Receiver().(starlark.HasSetKey) v, found, err := recv.Get(key) if err != nil { return starlark.None, fmt.Errorf("%s: %w", b.Name(), err) } if !found { v = dflt if err := recv.SetKey(key, dflt); err != nil { return starlark.None, fmt.Errorf("%s: %w", b.Name(), err) } } return v, nil } // https://github.com/google/starlark-go/blob/master/doc/spec.md#dict·update func dictUpdate(b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { // Unpack the arguments if len(args) > 1 { return nil, fmt.Errorf("update: got %d arguments, want at most 1", len(args)) } // Get the target dict := b.Receiver().(starlark.HasSetKey) if len(args) == 1 { switch updates := args[0].(type) { case starlark.IterableMapping: // Iterate over dict's key/value pairs, not just keys. for _, item := range updates.Items() { if err := dict.SetKey(item[0], item[1]); err != nil { return nil, err // dict is frozen } } default: // all other sequences iter := starlark.Iterate(updates) if iter == nil { return nil, fmt.Errorf("got %s, want iterable", updates.Type()) } defer iter.Done() var pair starlark.Value for i := 0; iter.Next(&pair); i++ { iterErr := func() error { iter2 := starlark.Iterate(pair) if iter2 == nil { return fmt.Errorf("dictionary update sequence element #%d is not iterable (%s)", i, pair.Type()) } defer iter2.Done() length := starlark.Len(pair) if length < 0 { return fmt.Errorf("dictionary update sequence element #%d has unknown length (%s)", i, pair.Type()) } else if length != 2 { return fmt.Errorf("dictionary update sequence element #%d has length %d, want 2", i, length) } var k, v starlark.Value iter2.Next(&k) iter2.Next(&v) return dict.SetKey(k, v) }() if iterErr != nil { return nil, iterErr } } } } // Then add the kwargs. before := starlark.Len(dict) for _, pair := range kwargs { if err := dict.SetKey(pair[0], pair[1]); err != nil { return nil, err // dict is frozen } } // In the common case, each kwarg will add another dict entry. // If that's not so, check whether it is because there was a duplicate kwarg. if starlark.Len(dict) < before+len(kwargs) { keys := make(map[starlark.String]bool, len(kwargs)) for _, kv := range kwargs { k := kv[0].(starlark.String) if keys[k] { return nil, fmt.Errorf("duplicate keyword arg: %v", k) } keys[k] = true } } return starlark.None, nil } // https://github.com/google/starlark-go/blob/master/doc/spec.md#dict·items func dictItems(b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { if err := starlark.UnpackPositionalArgs(b.Name(), args, kwargs, 0); err != nil { return starlark.None, fmt.Errorf("%s: %w", b.Name(), err) } items := b.Receiver().(starlark.IterableMapping).Items() res := make([]starlark.Value, 0, len(items)) for _, item := range items { res = append(res, item) // convert [2]starlark.Value to starlark.Value } return starlark.NewList(res), nil } // https://github.com/google/starlark-go/blob/master/doc/spec.md#dict·keys func dictKeys(b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { if err := starlark.UnpackPositionalArgs(b.Name(), args, kwargs, 0); err != nil { return starlark.None, fmt.Errorf("%s: %w", b.Name(), err) } items := b.Receiver().(starlark.IterableMapping).Items() res := make([]starlark.Value, 0, len(items)) for _, item := range items { res = append(res, item[0]) } return starlark.NewList(res), nil } // https://github.com/google/starlark-go/blob/master/doc/spec.md#dict·update func dictValues(b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { if err := starlark.UnpackPositionalArgs(b.Name(), args, kwargs, 0); err != nil { return starlark.None, fmt.Errorf("%s: %w", b.Name(), err) } items := b.Receiver().(starlark.IterableMapping).Items() res := make([]starlark.Value, 0, len(items)) for _, item := range items { res = append(res, item[1]) } return starlark.NewList(res), nil }