Merge pull request #375 from lcd1232/get_set

Support get/set

context.go

@@ -2,6 +2,7 @@ package telebot
 
 import (
 	"strings"
+	"sync"
 
 	"github.com/pkg/errors"
 )
@@ -119,6 +120,12 @@ type Context interface {
 	// Respond sends a response for the current callback query.
 	// See Respond from bot.go.
 	Respond(resp ...*CallbackResponse) error
+
+	// Get retrieves data from the context.
+	Get(key string) interface{}
+
+	// Set saves data in the context.
+	Set(key string, val interface{})
 }
 
 // nativeContext is a native implementation of the Context interface.
@@ -134,6 +141,9 @@ type nativeContext struct {
 	preCheckoutQuery   *PreCheckoutQuery
 	poll               *Poll
 	pollAnswer         *PollAnswer
+
+	lock  sync.RWMutex
+	store map[string]interface{}
 }
 
 func (c *nativeContext) Message() *Message {
@@ -349,3 +359,19 @@ func (c *nativeContext) Respond(resp ...*CallbackResponse) error {
 	}
 	return c.b.Respond(c.callback, resp...)
 }
+
+func (c *nativeContext) Set(key string, value interface{}) {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	if c.store == nil {
+		c.store = make(map[string]interface{})
+	}
+	c.store[key] = value
+}
+
+func (c *nativeContext) Get(key string) interface{} {
+	c.lock.RLock()
+	defer c.lock.RUnlock()
+	return c.store[key]
+}

context_test.go

@@ -1,3 +1,16 @@
 package telebot
 
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
 var _ Context = (*nativeContext)(nil)
+
+func TestContextStore(t *testing.T) {
+	var c Context
+	c = new(nativeContext)
+	c.Set("name", "Jon Snow")
+	assert.Equal(t, "Jon Snow", c.Get("name"))
+}