Allow a different subtotal function for each data field.

2 files changed, 42 insertions, 26 deletions

diff --git a/lib/axlsx/workbook/worksheet/pivot_table.rb b/lib/axlsx/workbook/worksheet/pivot_table.rb
index baa08988..da87162d 100644
--- a/lib/axlsx/workbook/worksheet/pivot_table.rb
+++ b/lib/axlsx/workbook/worksheet/pivot_table.rb
@@ -23,13 +23,10 @@ module Axlsx
       @columns = []
       @data = []
       @pages = []
-      @subtotal = 'sum'
       parse_options options
       yield self if block_given?
     end
 
-    attr_writer :subtotal
-
     # The reference to the table data
     # @return [String]
     attr_reader :ref
@@ -88,10 +85,17 @@ module Axlsx
     # (see #data)
     def data=(v)
       DataTypeValidator.validate "#{self.class}.data", [Array], v
-      v.each do |ref|
-        DataTypeValidator.validate "#{self.class}.data[]", [String], ref
+      @data = []
+      v.each do |data_field|
+        if data_field.is_a? String
+          data_field = {:ref => data_field}
+        end
+        data_field.values.each do |value|
+          DataTypeValidator.validate "#{self.class}.data[]", [String], value
+        end
+        @data << data_field
       end
-      @data = v
+      @data
     end
 
     # The pages
@@ -189,11 +193,11 @@ module Axlsx
         str << '</pageFields>'
       end
       unless data.empty?
-        str << '<dataFields count="' << data.size.to_s << '">'
+        str << "<dataFields count=\"#{data.size}\">"
         data.each do |datum_value|
-          str << "<dataField name=\"#{@subtotal} of " << datum_value << '" ' <<
-                            'fld="' << header_index_of(datum_value).to_s << '" ' <<
-                            'baseField="0" baseItem="0" subtotal="' << @subtotal << '"/>'
+          str << "<dataField name='#{@subtotal} of #{datum_value[:ref]}' fld='#{header_index_of(datum_value[:ref])}' baseField='0' baseItem='0'"
+          str << " subtotal='#{datum_value[:subtotal]}' " if datum_value[:subtotal]
+          str << "/>"
         end
         str << '</dataFields>'
       end
diff --git a/test/workbook/worksheet/tc_pivot_table.rb b/test/workbook/worksheet/tc_pivot_table.rb
index 9adbaf93..d909eeb6 100644
--- a/test/workbook/worksheet/tc_pivot_table.rb
+++ b/test/workbook/worksheet/tc_pivot_table.rb
@@ -1,5 +1,17 @@
 require 'tc_helper.rb'
 
+
+def shared_test_pivot_table_xml_validity(pivot_table)
+  schema = Nokogiri::XML::Schema(File.open(Axlsx::SML_XSD))
+  doc = Nokogiri::XML(pivot_table.to_xml_string)
+  errors = []
+  schema.validate(doc).each do |error|
+    errors.push error
+    puts error.message
+  end
+  assert(errors.empty?, "error free validation")
+end
+
 class TestPivotTable < Test::Unit::TestCase
   def setup
     p = Axlsx::Package.new
@@ -33,10 +45,17 @@ class TestPivotTable < Test::Unit::TestCase
     end
     assert_equal(['Year', 'Month'], pivot_table.rows)
     assert_equal(['Type'], pivot_table.columns)
-    assert_equal(['Sales'], pivot_table.data)
+    assert_equal([{:ref=>"Sales"}], pivot_table.data)
     assert_equal(['Region'], pivot_table.pages)
   end
 
+  def test_add_pivot_table_with_options_on_data_field
+    pivot_table = @ws.add_pivot_table('G5:G6', 'A1:D5') do |pt|
+      pt.data = [{:ref=>"Sales", :subtotal => 'average'}]
+    end
+    assert_equal([{:ref=>"Sales", :subtotal => 'average'}], pivot_table.data)
+  end
+
   def test_header_indices
     pivot_table = @ws.add_pivot_table('G5:G6', 'A1:E5')
     assert_equal(0,   pivot_table.header_index_of('Year'   ))
@@ -73,14 +92,7 @@ class TestPivotTable < Test::Unit::TestCase
 
   def test_to_xml_string
     pivot_table = @ws.add_pivot_table('G5:G6', 'A1:D5')
-    schema = Nokogiri::XML::Schema(File.open(Axlsx::SML_XSD))
-    doc = Nokogiri::XML(pivot_table.to_xml_string)
-    errors = []
-    schema.validate(doc).each do |error|
-      errors.push error
-      puts error.message
-    end
-    assert(errors.empty?, "error free validation")
+    shared_test_pivot_table_xml_validity(pivot_table)
   end
 
   def test_to_xml_string_with_configuration
@@ -90,13 +102,13 @@ class TestPivotTable < Test::Unit::TestCase
       pt.data = ['Sales']
       pt.pages = ['Region']
     end
-    schema = Nokogiri::XML::Schema(File.open(Axlsx::SML_XSD))
-    doc = Nokogiri::XML(pivot_table.to_xml_string)
-    errors = []
-    schema.validate(doc).each do |error|
-      errors.push error
-      puts error.message
+    shared_test_pivot_table_xml_validity(pivot_table)
+  end
+
+  def test_to_xml_string_with_options_on_data_field
+    pivot_table = @ws.add_pivot_table('G5:G6', 'A1:E5') do |pt|
+      pt.data = [{:ref=>"Sales", :subtotal => 'average'}]
     end
-    assert(errors.empty?, "error free validation")
+    shared_test_pivot_table_xml_validity(pivot_table)
   end
 end