lib/junoser/display/set.rb in junoser-0.1.2 vs lib/junoser/display/set.rb in junoser-0.1.3

@@ -1,17 +1,23 @@
 require 'parslet'
-require 'junoser/display/base'
+require 'junoser/input'
 
 module Junoser
   module Display
     class Set
-      include Base
+      def initialize(io_or_string)
+        @input = io_or_string
+      end
 
       def transform
+        result = ''
+
         process do |current_stack, str|
-          @output.puts transform_line(current_stack, str)
+          result << transform_line(current_stack, str) << "\n"
         end
+
+        result
       end
 
       def commit_check(&block)
         parser = Junoser::Parser.new
         passed = true
@@ -38,10 +44,10 @@
       private
 
       def process(&block)
         stack = []
 
-        read_io_or_string.split("\n").each do |line|
+        Junoser::Input.new(@input).read.split("\n").each do |line|
           case line
           when /(.*){/
             stack.push $1.strip
           when '}'
             stack.pop