From 0138fd39e955c08f0c4d2e246532c2725d3b1021 Mon Sep 17 00:00:00 2001
From: Eloy Espinaco <eloyesp@gmail.com>
Date: Sun, 11 Dec 2011 11:52:00 -0300
Subject: [PATCH] Adding compatibility for different encodings - fix #31

I change the default internal encoding because seed.inspect need to be
properly encoded.
---
 lib/seed-fu/writer.rb | 24 +++++++++++++++++++++++-
 spec/writer_spec.rb   | 14 ++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/lib/seed-fu/writer.rb b/lib/seed-fu/writer.rb
index 04e28fd..bfe1c25 100644
--- a/lib/seed-fu/writer.rb
+++ b/lib/seed-fu/writer.rb
@@ -31,9 +31,11 @@ class Writer
     # @option options [:seed, :seed_once] :seed_type (:seed) The method to use when generating
     #   seeds. See {ActiveRecordExtension} for details.
     # @option options [Array<Symbol>] :constraints ([:id]) The constraining attributes for the seeds
+    # @option options [String] :encoding The encoding to use in the generated file.
     def initialize(options = {})
       @options = self.class.default_options.merge(options)
       raise ArgumentError, "missing option :class_name" unless @options[:class_name]
+      set_encoding if @options[:encoding]
     end
 
     # Creates a new instance of {Writer} with the `options`, and then calls {#write} with the
@@ -56,10 +58,11 @@ def write(io_or_filename, &block)
       if io_or_filename.respond_to?(:write)
         write_to_io(io_or_filename, &block)
       else
-        File.open(io_or_filename, 'w') do |file|
+        File.open(io_or_filename, "w#{encoding_string}") do |file|
           write_to_io(file, &block)
         end
       end
+      unset_encoding if @old_encoding
     end
 
     # Add a seed. Must be called within a block passed to {#write}.
@@ -88,6 +91,7 @@ def <<(seed)
 
       def write_to_io(io)
         @io, @count = io, 0
+        @io.write(encoding_comment) if @options[:encoding]
         @io.write(file_header)
         @io.write(seed_header)
         yield(self)
@@ -128,5 +132,23 @@ def seed_footer
       def chunk_this_seed?
         @count != 0 && (@count % @options[:chunk_size]) == 0
       end
+
+      def encoding_string
+        ":#{@options[:encoding]}" if @options[:encoding]
+      end
+
+      def encoding_comment
+        "# encoding: #{@options[:encoding]}\n"
+      end
+
+      def set_encoding
+        @old_encoding = Encoding.default_internal
+        Encoding.default_internal = @options[:encoding]
+      end
+
+      def unset_encoding
+        Encoding.default_internal = @old_encoding if @old_encoding
+      end
   end
 end
+
diff --git a/spec/writer_spec.rb b/spec/writer_spec.rb
index 24eb7ed..32bf7f4 100644
--- a/spec/writer_spec.rb
+++ b/spec/writer_spec.rb
@@ -1,3 +1,4 @@
+# encoding: utf-8
 require 'spec_helper'
 
 describe SeedFu::Writer do
@@ -42,4 +43,17 @@
 
     SeededModel.find(1).title.should == "Dr"
   end
+
+  it "should support specifying the output encoding to use" do
+    SeedFu::Writer.write(@file_name, :class_name => 'SeededModel', :encoding => "utf-8") do |writer|
+      writer << { :id => 1, :title => "Mr" }
+      writer << { :id => 2, :title => "Máster" }
+    end
+
+    File.read(@file_name).should include("# encoding: utf-8\n")
+
+    load @file_name
+    SeededModel.find(2).title.should == "Máster"
+  end
 end
+