Re-add CurrentContext.with

Author: binarylogic

lib/timber/current_context.rb

@@ -38,6 +38,11 @@ def remove(*args)
       def reset(*args)
         instance.reset(*args)
       end
+
+      # Convenience method for {CurrentContext#with}. See {CurrentContext#with} for more info.
+      def with(*args, &block)
+        instance.with(*args, &block)
+      end
     end
 
     # Adds contexts but does not remove them. See {#with} for automatic maintenance and {#remove}
@@ -68,6 +73,12 @@ def remove(*keys)
       self
     end
 
+    def replace(hash)
+      @hash = hash
+      expire_cache!
+      self
+    end
+
     # Resets the context to be blank. Use this carefully! This will remove *any* context,
     # include context that is automatically included with Timber.
     def reset
@@ -83,6 +94,35 @@ def snapshot
       @snapshot ||= hash.clone
     end
 
+    # Adds a context and then removes it when the block is finished executing.
+    #
+    # @note Because context is included with every log line, it is recommended that you limit this
+    #   to only neccessary data.
+    #
+    # @example Adding a custom context
+    #   Timber::CurrentContext.with({build: {version: "1.0.0"}}) do
+    #     # ... anything logged here will include the context ...
+    #   end
+    #
+    # @note Any custom context needs to have a single root key to be valid. i.e. instead of:
+    #   Timber::CurrentContext.with(job_id: "123", job_name: "Refresh User Account")
+    #
+    # do
+    #
+    #   Timber::CurrentContext.with(job: {job_id: "123", job_name: "Refresh User Account"})
+    #
+    # @example Adding multiple contexts
+    #   Timber::CurrentContext.with(context1, context2) { ... }
+    def with(*objects)
+      old_hash = hash.clone
+      begin
+        add(*objects)
+        yield
+      ensure
+        replace(old_hash)
+      end
+    end
+
     private
       # The internal hash that is maintained. Use {#with} and {#add} for hash maintenance.
       def hash

spec/timber/current_context_spec.rb

@@ -89,4 +89,25 @@
       expect(described_class.instance.send(:hash)[:build]).to be_nil
     end
   end
+
+  describe ".with" do
+    it "should merge the context and cleanup on block exit" do
+      expect(described_class.instance.send(:hash)[:build]).to be_nil
+
+      described_class.with({build: {version: "1.0.0"}}) do
+        expect(described_class.instance.send(:hash)[:build]).to eq({:version=>"1.0.0"})
+
+        described_class.with({testing: {key: "value"}}) do
+          expect(described_class.instance.send(:hash)[:build]).to eq({:version=>"1.0.0"})
+          expect(described_class.instance.send(:hash)[:testing]).to eq({:key=>"value"})
+        end
+
+        expect(described_class.instance.send(:hash)[:build]).to eq({:version=>"1.0.0"})
+        expect(described_class.instance.send(:hash)[:testing]).to be_nil
+      end
+
+      expect(described_class.instance.send(:hash)[:build]).to be_nil
+      expect(described_class.instance.send(:hash)[:testing]).to be_nil
+    end
+  end
 end