diff --git a/mctx/annotate.go b/mctx/annotate.go new file mode 100644 index 0000000..b2d72e9 --- /dev/null +++ b/mctx/annotate.go @@ -0,0 +1,71 @@ +package mctx + +import ( + "context" + "fmt" +) + +type annotateKey struct { + userKey interface{} +} + +// Annotate takes in one or more key/value pairs (kvs' length must be even) and +// returns a Context carrying them. Annotations only exist on the local level, +// i.e. a child and parent share different annotation namespaces. +func Annotate(ctx context.Context, kvs ...interface{}) context.Context { + for i := 0; i < len(kvs); i += 2 { + ctx = WithLocalValue(ctx, annotateKey{kvs[i]}, kvs[i+1]) + } + return ctx +} + +// Annotations describes a set of keys/values which were set on a Context (but +// not its parents or children) using Annotate. +type Annotations [][2]interface{} + +// LocalAnnotations returns all key/value pairs which have been set via Annotate +// on this Context (but not its parent or children). If a key was set twice then +// only the most recent value is included. The returned order is +// non-deterministic. +func LocalAnnotations(ctx context.Context) Annotations { + var annotations Annotations + localValuesIter(ctx, func(key, val interface{}) { + aKey, ok := key.(annotateKey) + if !ok { + return + } + annotations = append(annotations, [2]interface{}{aKey.userKey, val}) + }) + return annotations +} + +// StringMap formats each of the key/value pairs into strings using fmt.Sprint. +// If any two keys format to the same string, then type information will be +// prefaced to each one. +func (aa Annotations) StringMap() map[string]string { + keyTypes := make(map[string]interface{}, len(aa)) + keyVals := map[string]string{} + setKey := func(k, v interface{}) { + kStr := fmt.Sprint(k) + oldType := keyTypes[kStr] + if oldType == nil { + keyTypes[kStr] = k + keyVals[kStr] = fmt.Sprint(v) + return + } + + // check if oldKey is in kvm, if so it needs to be moved to account for + // its type info + if oldV, ok := keyVals[kStr]; ok { + delete(keyVals, kStr) + keyVals[fmt.Sprintf("%T(%s)", oldType, kStr)] = oldV + } + + keyVals[fmt.Sprintf("%T(%s)", k, kStr)] = fmt.Sprint(v) + } + + for _, kv := range aa { + setKey(kv[0], kv[1]) + } + return keyVals +} diff --git a/mctx/annotate_test.go b/mctx/annotate_test.go new file mode 100644 index 0000000..3db3c04 --- /dev/null +++ b/mctx/annotate_test.go @@ -0,0 +1,55 @@ +package mctx + +import ( + "context" + . "testing" + + "github.com/mediocregopher/mediocre-go-lib/mtest/massert" +) + +func TestAnnotate(t *T) { + parent := context.Background() + parent = Annotate(parent, "a", "foo") + parent = Annotate(parent, "b", "bar") + + child := NewChild(parent, "child") + child = Annotate(child, "a", "FOO") + child = Annotate(child, "c", "BAZ") + parent = WithChild(parent, child) + + parentAnnotations := LocalAnnotations(parent) + childAnnotations := LocalAnnotations(child) + massert.Fatal(t, massert.All( + massert.Len(parentAnnotations, 2), + massert.Has(parentAnnotations, [2]interface{}{"a", "foo"}), + massert.Has(parentAnnotations, [2]interface{}{"b", "bar"}), + massert.Len(childAnnotations, 2), + massert.Has(childAnnotations, [2]interface{}{"a", "FOO"}), + massert.Has(childAnnotations, [2]interface{}{"c", "BAZ"}), + )) +} + +func TestAnnotationsStingMap(t *T) { + type A int + type B int + aa := Annotations{ + {"foo", "bar"}, + {"1", "one"}, + {1, 1}, + {0, 0}, + {A(0), 0}, + {B(0), 0}, + } + + err := massert.Equal(map[string]string{ + "foo": "bar", + "string(1)": "one", + "int(1)": "1", + "int(0)": "0", + "mctx.A(0)": "0", + "mctx.B(0)": "0", + }, aa.StringMap()).Assert() + if err != nil { + t.Fatal(err) + } +} diff --git a/mctx/ctx.go b/mctx/ctx.go index de4e3a8..304163f 100644 --- a/mctx/ctx.go +++ b/mctx/ctx.go @@ -218,17 +218,26 @@ func LocalValue(ctx context.Context, key interface{}) interface{} { } } -// LocalValues returns all key/value pairs which have been set on the Context -// via WithLocalValue. -func LocalValues(ctx context.Context) map[interface{}]interface{} { - m := map[interface{}]interface{}{} +func localValuesIter(ctx context.Context, callback func(key, val interface{})) { + m := map[interface{}]struct{}{} lv, _ := ctx.Value(localValsKey(0)).(*localVal) for { if lv == nil { - return m + return } else if _, ok := m[lv.key]; !ok { - m[lv.key] = lv.val + callback(lv.key, lv.val) + m[lv.key] = struct{}{} } lv = lv.prev } } + +// LocalValues returns all key/value pairs which have been set on the Context +// via WithLocalValue. +func LocalValues(ctx context.Context) map[interface{}]interface{} { + m := map[interface{}]interface{}{} + localValuesIter(ctx, func(key, val interface{}) { + m[key] = val + }) + return m +}