From 0b4856b7417592743a27fc3e30867fd141395bf6 Mon Sep 17 00:00:00 2001
From: Rimas Silkaitis <neovintage@gmail.com>
Date: Thu, 6 Apr 2017 11:43:41 -0700
Subject: [PATCH] User defined context store types (#339)

allow the context storage to handle any types
---
 spec/context_spec.cr        |  7 +++++++
 spec/spec_helper.cr         | 17 +++++++++++++++++
 src/kemal/ext/context.cr    |  9 +++++++--
 src/kemal/helpers/macros.cr | 12 ++++++++++++
 4 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/spec/context_spec.cr b/spec/context_spec.cr
index b96101a..67f364f 100644
--- a/spec/context_spec.cr
+++ b/spec/context_spec.cr
@@ -34,9 +34,14 @@ describe "Context" do
 
   it "can store variables" do
     before_get "/" do |env|
+      t = TestContextStorageType.new
+      t.id = 32
+      a = AnotherContextStorageType.new
       env.set "key", "value"
       env.set "before_get", "Kemal"
       env.set "before_get_int", 123
+      env.set "before_get_context_test", t
+      env.set "another_context_test", a
       env.set "before_get_float", 3.5
     end
 
@@ -47,6 +52,7 @@ describe "Context" do
         before_get:       env.get("before_get"),
         before_get_int:   env.get("before_get_int"),
         before_get_float: env.get("before_get_float"),
+        before_get_context_test: env.get("before_get_context_test"),
       }
     end
     request = HTTP::Request.new("GET", "/")
@@ -59,5 +65,6 @@ describe "Context" do
     context.store["before_get"].should eq "Kemal"
     context.store["before_get_int"].should eq 123
     context.store["before_get_float"].should eq 3.5
+    context.store["before_get_context_test"].as(TestContextStorageType).id.should eq 32
   end
 end
diff --git a/spec/spec_helper.cr b/spec/spec_helper.cr
index 9c2b488..bdd5790 100644
--- a/spec/spec_helper.cr
+++ b/spec/spec_helper.cr
@@ -12,6 +12,23 @@ class CustomLogHandler < Kemal::BaseLogHandler
   end
 end
 
+class TestContextStorageType
+  property id
+  @id = 1
+
+  def to_s
+    @id
+  end
+end
+
+class AnotherContextStorageType
+  property name
+  @name = "kemal-context"
+end
+
+add_context_storage_type(TestContextStorageType)
+add_context_storage_type(AnotherContextStorageType)
+
 def create_request_and_return_io(handler, request)
   io = IO::Memory.new
   response = HTTP::Server::Response.new(io)
diff --git a/src/kemal/ext/context.cr b/src/kemal/ext/context.cr
index 0ca9af7..cd3a756 100644
--- a/src/kemal/ext/context.cr
+++ b/src/kemal/ext/context.cr
@@ -4,8 +4,13 @@
 # Instances of this class are passed to an `HTTP::Server` handler.
 class HTTP::Server
   class Context
-    alias StoreTypes = Nil | String | Int32 | Int64 | Float64 | Bool
-    getter store = {} of String => StoreTypes
+    # :nodoc:
+    STORE_MAPPINGS = [ Nil, String, Int32, Int64, Float64, Bool ]
+
+    macro finished
+      alias StoreTypes = Union({{ *STORE_MAPPINGS }})
+      getter store = {} of String => StoreTypes
+    end
 
     def params
       @request.url_params ||= route_lookup.params
diff --git a/src/kemal/helpers/macros.cr b/src/kemal/helpers/macros.cr
index a69fb46..e217197 100644
--- a/src/kemal/helpers/macros.cr
+++ b/src/kemal/helpers/macros.cr
@@ -76,3 +76,15 @@ macro halt(env, status_code = 200, response = "")
   {{env}}.response.close
   next
 end
+
+# Extends context storage with user defined types.
+#
+# class User
+#   property name
+# end
+#
+# add_context_storage_type(User)
+#
+macro add_context_storage_type(type)
+  {{ HTTP::Server::Context::STORE_MAPPINGS.push(type) }}
+end