mctx: implement beginnings of Annotations functionality

This commit is contained in:
Brian Picciano 2019-02-08 18:44:20 -05:00
parent 4b446a0efc
commit 000a57689d
3 changed files with 141 additions and 6 deletions

71
mctx/annotate.go Normal file
View File

@ -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
}

55
mctx/annotate_test.go Normal file
View File

@ -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)
}
}

View File

@ -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
}