From a31f07f7c22f9b85d1a7b6a3ef5fdeb009cb83c0 Mon Sep 17 00:00:00 2001
From: Holger Just <hello@holgerjust.de>
Date: Thu, 21 Sep 2017 15:52:56 +0200
Subject: [PATCH] Yield the last created flow on Logger.new

This allows to configure the last (and often only) flow on creation of
the Logger without having to manually create the object. We can thus use
the following shortcut to e.g. set a custom encoder:

    Rackstash::Logger.new(STDOUT) do
      encoder Rackstash::Encoder::Message.new
    end
---
 lib/rackstash/logger.rb       | 13 ++++++++++++-
 spec/rackstash/logger_spec.rb | 29 +++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/lib/rackstash/logger.rb b/lib/rackstash/logger.rb
index ac9636f..4a51425 100644
--- a/lib/rackstash/logger.rb
+++ b/lib/rackstash/logger.rb
@@ -57,13 +57,24 @@ module Rackstash
     #   By default we use {PROGNAME}.
     # @param formatter [#call] the log formatter for each individual buffered
     #   line. See {#formatter} for details.
-    def initialize(*flows, level: DEBUG, progname: PROGNAME, formatter: Formatter.new)
+    # @yieldparam flow [Rackstash::Flow] if the given block accepts an argument,
+    #   we yield the last {Flow} as a parameter. Without an expected argument,
+    #   the block is directly executed in the context of the last {Flow}.
+    def initialize(*flows, level: DEBUG, progname: PROGNAME, formatter: Formatter.new, &block)
       @buffer_stack = Concurrent::ThreadLocalVar.new
 
       @sink = Rackstash::Sink.new(*flows)
       self.level = level
       self.progname = progname
       self.formatter = formatter
+
+      if block_given? && (flow = @sink.flows.last)
+        if block.arity == 0
+          flow.instance_eval(&block)
+        else
+          yield flow
+        end
+      end
     end
 
     # Add a message to the current buffer without any further formatting. If
diff --git a/spec/rackstash/logger_spec.rb b/spec/rackstash/logger_spec.rb
index d008434..a523451 100644
--- a/spec/rackstash/logger_spec.rb
+++ b/spec/rackstash/logger_spec.rb
@@ -35,6 +35,35 @@ describe Rackstash::Logger do
       logger = described_class.new('output.log', formatter: -> {})
       expect(logger.formatter).to be_a Proc
     end
+
+    it 'yields the last flow to a parameterized block' do
+      block_called = 0
+      block_self = nil
+      block_args = nil
+
+      described_class.new(StringIO.new, StringIO.new) do |*args|
+        block_called += 1
+        block_self = self
+        block_args = args
+      end
+
+      expect(block_called).to eql 1
+      expect(block_self).to equal self
+      expect(block_args).to match [instance_of(Rackstash::Flow)]
+    end
+
+    it 'instance_evals the parameter-less block in the last flow' do
+      block_called = 0
+      block_self = nil
+
+      described_class.new(StringIO.new, StringIO.new) do
+        block_called += 1
+        block_self = self
+      end
+
+      expect(block_called).to eql 1
+      expect(block_self).to be_instance_of(Rackstash::Flow)
+    end
   end
 
   describe 'subscript accessors' do