lib/timber/current_context.rb in timber-3.0.0 vs lib/timber/current_context.rb in timber-3.0.1

@@ -36,10 +36,15 @@
 
       # Convenience method for {CurrentContext#reset}. See {CurrentContext#reset} for more info.
       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}
     # to remove them yourself.
     #
@@ -66,10 +71,16 @@
       end
       expire_cache!
       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
       hash.clear
       expire_cache!
@@ -79,9 +90,38 @@
     # Snapshots the current context so that you get a moment in time representation of the context,
     # since the context can change as execution proceeds. Note that individual contexts
     # should be immutable, and we implement snapshot caching as a result of this assumption.
     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