lib/timber/current_context.rb in timber-1.1.13 vs lib/timber/current_context.rb in timber-1.1.14

@@ -21,17 +21,21 @@
       # Convenience method for {#add}. See {#add} for full details and examples.
       def add(*args)
         instance.add(*args)
       end
 
+      def hash(*args)
+        instance.hash(*args)
+      end
+
       # Convenience method for {#remove}. See {#remove} for full details and examples.
       def remove(*args)
         instance.remove(*args)
       end
 
-      def hash(*args)
-        instance.hash(*args)
+      def reset(*args)
+        instance.reset(*args)
       end
     end
 
     # Adds a context and then removes it when the block is finished executing.
     #
@@ -54,22 +58,15 @@
     #   # Be sure to checkout Timber::Contexts! These are officially supported and many of these
     #   # will be automatically included via Timber::Probes
     #
     # @example Adding multiple contexts
     #   Timber::CurrentContext.with(context1, context2) { ... }
-    def with(*contexts)
-      add(*contexts)
+    def with(*objects)
+      add(*objects)
       yield
     ensure
-      contexts.each do |context|
-        if context.keyspace == :custom
-          # Custom contexts are merged and should be removed the same
-          hash[context.keyspace].delete(context.type)
-        else
-          remove(context)
-        end
-      end
+      remove(*objects)
     end
 
     # Adds contexts but does not remove them. See {#with} for automatic maintenance and {#remove}
     # to remove them yourself.
     #
@@ -81,33 +78,53 @@
         key = context.keyspace
         json = context.as_json # Convert to json now so that we aren't doing it for every line
         if key == :custom
           # Custom contexts are merged into the space
           hash[key] ||= {}
-          hash[key].merge(json)
+          hash[key].merge!(json)
         else
           hash[key] = json
         end
       end
     end
 
-    # Removes a context. This must be a {Timber::Context} type. See {Timber::Contexts} for a list.
-    # If you wish to remove by key, or some other way, use {#hash} and modify the hash accordingly.
-    def remove(*contexts)
-      contexts.each do |context|
-        hash.delete(context.keyspace)
-      end
-    end
-
     # The internal hash that is maintained. It is recommended that you use {#with} and {#add}
     # for hash maintenance.
     def hash
       Thread.current[THREAD_NAMESPACE] ||= {}
     end
 
+    # Removes a context. If you wish to remove by key, or some other way, use {#hash} and
+    # modify the hash accordingly.
+    def remove(*objects)
+      objects.each do |object|
+        if object.is_a?(Symbol)
+          hash.delete(object)
+        else
+          context = Contexts.build(object)
+
+          if context.keyspace == :custom
+            # Custom contexts are merged and should be removed the same
+            hash[context.keyspace].delete(context.type)
+            if hash[context.keyspace] == {}
+              hash.delete(context.keyspace)
+            end
+          else
+            hash.delete(context.keyspace)
+          end
+        end
+      end
+    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
+    end
+
     # Snapshots the current context so that you get a moment in time representation of the context,
     # since the context can change as execution proceeds.
     def snapshot
       hash.clone
     end
   end
-end
\ No newline at end of file
+end