Finish tty spec and simplified a bit more

.gitignore

@@ -1 +1,4 @@
-pkg
+pkg
+*.swp
+extensions
+t.rb

bin/dietrb

@@ -32,9 +32,9 @@ end
 IRB.formatter.filter_from_backtrace << /^#{__FILE__}/
 
 if ARGV.empty?
-  # require 'irb/driver/readline'
+  require 'irb/driver/readline'
   # require 'irb/driver/tty'
-  require 'irb/driver/socket'
+  # require 'irb/driver/socket'
 else
   path = ARGV.shift
   require 'irb/driver/file'

lib/irb/driver/tty.rb

@@ -23,22 +23,15 @@ def consume(context)
         ""
       end
 
+      # Feeds input into a given context.
+      #
+      # Ensures that the standard output object is a OutputRedirector, or a
+      # subclass thereof.
       def run(context)
-        ensure_output_redirector do
-          while line = consume(context)
-            continue = context.process_line(line)
-            break unless continue
-          end
+        before, $stdout = $stdout, OutputRedirector.new unless $stdout.is_a?(OutputRedirector)
+        while line = consume(context)
+          break unless context.process_line(line)
         end
-      end
-
-      # Ensure that the standard output object is a OutputRedirector. If it's
-      # already a OutputRedirector, do nothing.
-      def ensure_output_redirector
-        unless $stdout.is_a?(IRB::Driver::OutputRedirector)
-          before, $stdout = $stdout, IRB::Driver::OutputRedirector.new
-        end
-        yield
       ensure
         $stdout = before if before
       end
@@ -55,4 +48,4 @@ def irb(object, binding = nil)
   end
 
   private :irb
-end
+end

spec/driver/tty_spec.rb

@@ -24,4 +24,20 @@ def @driver.readline(_); raise Interrupt; end
     @driver.consume(@context).should == ""
     @context.source.to_s.should == ""
   end
-end
+
+  it "feeds input into a given context" do
+    $from_context = false
+    @driver.input.stub_input "$from_context = true", "exit"
+    @driver.run(@context)
+    $from_context.should == true
+  end
+
+  it "makes sure there's one global output redirector while running a context" do
+    before = $stdout
+    $from_context = nil
+    @driver.input.stub_input "$from_context = $stdout", "exit"
+    @driver.run(@context)
+    $from_context.class == IRB::Driver::OutputRedirector
+    $stdout.should == before
+  end
+end