Skip to content

Commit

Permalink
add generic context utility (#41)
Browse files Browse the repository at this point in the history
  • Loading branch information
matoszz authored Dec 8, 2024
1 parent 78571ec commit 26bbb15
Show file tree
Hide file tree
Showing 3 changed files with 164 additions and 0 deletions.
54 changes: 54 additions & 0 deletions contextx/contextx.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
package contextx

import (
"context"
)

// key is a unique type that we can use as a key in a context
type key[T any] struct{}

// With returns a copy of parent that contains the given value which can be retrieved by calling From with the resulting context
// The function uses a generic key type to ensure that the stored value is type-safe and can be uniquely identified and retrieved without
// risk of key collisions
func With[T any](ctx context.Context, v T) context.Context {
return context.WithValue(ctx, key[T]{}, v)
}

// From returns the value associated with the wanted type from the context
// It performs a type assertion to convert the value to the desired type T
// If the type assertion is successful, it returns the value and true
// If the type assertion fails, it returns the zero value of type T and false
func From[T any](ctx context.Context) (T, bool) {
v, ok := ctx.Value(key[T]{}).(T)

return v, ok
}

// MustFrom is similar to from, except that it panics if the type assertion fails / the value is not in the context
func MustFrom[T any](ctx context.Context) T {
return ctx.Value(key[T]{}).(T)
}

// FromOr returns the value associated with the wanted type or the given default value if the type is not found
// This function is useful when you want to ensure that a value is always returned from the context, even if the
// context does not contain a value of the desired type. By providing a default value, you can avoid handling
// the case where the value is missing and ensure that your code has a fallback value to use
func FromOr[T any](ctx context.Context, def T) T {
v, ok := From[T](ctx)
if !ok {
return def
}

return v
}

// FromOrFunc returns the value associated with the wanted type or the result of the given function if the type is not found
// This function is useful when the default value is expensive to compute or when the default value depends on some runtime conditions
func FromOrFunc[T any](ctx context.Context, f func() T) T {
v, ok := From[T](ctx)
if !ok {
return f()
}

return v
}
102 changes: 102 additions & 0 deletions contextx/contextx_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,102 @@
package contextx

import (
"context"
"reflect"
"testing"
)

func TestNormalOperation(t *testing.T) {
ctx := context.Background()
ctx = With(ctx, 10)

if MustFrom[int](ctx) != 10 {
t.FailNow()
}

if _, ok := From[float64](ctx); ok {
t.FailNow()
}
}

func TestIsolatedFromExplicitTypeReflection(t *testing.T) {
ctx := context.Background()

ctx = With(ctx, 10)

ctx = context.WithValue(ctx, reflect.TypeOf(20), 20)

if MustFrom[int](ctx) != 10 {
t.FailNow()
}
}

func TestPanicIfNoValue(t *testing.T) {
defer func() {
if recover() == nil {
t.FailNow()
}
}()

MustFrom[int](context.Background())
}

type x interface {
a()
}

type y struct{ v int }

func (y) a() {}

type z struct{ f func() }

func (z z) a() { z.f() }

func TestShouldWorkOnInterface(t *testing.T) {
var a x = y{10}

ctx := context.Background()
ctx = With(ctx, a)

b := MustFrom[x](ctx)
if b.(y).v != 10 {
t.FailNow()
}

r := ""
a = z{func() { r = "hello" }}

ctx = With(ctx, a)

MustFrom[x](ctx).a()

if r != "hello" {
t.FailNow()
}
}
func TestFromOr(t *testing.T) {
ctx := context.Background()
ctx = With(ctx, 10)

if FromOr(ctx, 20) != 10 {
t.FailNow()
}

if FromOr(context.Background(), 20) != 20 {
t.FailNow()
}
}

func TestFromOrFunc(t *testing.T) {
ctx := context.Background()
ctx = With(ctx, 10)

if FromOrFunc(ctx, func() int { return 20 }) != 10 {
t.FailNow()
}

if FromOrFunc(context.Background(), func() int { return 20 }) != 20 {
t.FailNow()
}
}
8 changes: 8 additions & 0 deletions contextx/doc.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
// Package contextx is a helper package for managing context values
// Most **request-scoped data** is a singleton per request
// That is, it doesn't make sense for a request to carry around multiple loggers, users, traces
// you want to carry the _same one_ with you from function call to function call
// the way we've handled this historically is a separate context key per type you want to carry in the struct
// but with generics, instead of having to make a new zero-sized type for every struct
// we can just make a single generic type and use it for everything which is what this helper package is intended to do
package contextx

0 comments on commit 26bbb15

Please sign in to comment.