Add pivot table option to sort headers (#143)

5 files changed, 57 insertions, 14 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 306821a2..cf355f8f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,7 @@ CHANGELOG
 ---------
 
 - **Unreleased**
+  - [PR #143](https://github.com/caxlsx/caxlsx/pull/143) - Add setting `sort_on_headers` for pivot tables
   - [PR #132](https://github.com/caxlsx/caxlsx/pull/132) - Remove monkey patch from Object#instance_values
   - [PR #139](https://github.com/caxlsx/caxlsx/pull/139) - Sort archive entries for correct MIME detection with `file` command
   - [PR #140](https://github.com/caxlsx/caxlsx/pull/140) - Update gemspec to recent styles - it reduced the size of the gem
diff --git a/examples/images/pivot_table_example.png b/examples/images/pivot_table_example.png
index 70a625f8..80631748 100644
--- a/examples/images/pivot_table_example.png
+++ b/examples/images/pivot_table_example.png
diff --git a/examples/pivot_table_example.md b/examples/pivot_table_example.md
index 8dd12955..50d20c5f 100644
--- a/examples/pivot_table_example.md
+++ b/examples/pivot_table_example.md
@@ -24,8 +24,8 @@ wb.add_worksheet(name: 'Basic Worksheet') do |sheet|
     ]
   end
 
-  sheet.add_pivot_table 'G4:L17', 'A1:E31' do |pivot_table|
-    pivot_table.rows = ['Month', 'Year']
+  sheet.add_pivot_table 'G4:L17', 'A1:E31', sort_on_headers: ['Year', 'Month'] do |pivot_table|
+    pivot_table.rows = ['Year', 'Month']
     pivot_table.columns = ['Type']
     pivot_table.data = [ref: 'Sales', num_fmt: 4]
     pivot_table.pages = ['Region']
diff --git a/lib/axlsx/workbook/worksheet/pivot_table.rb b/lib/axlsx/workbook/worksheet/pivot_table.rb
index bc500a8a..60960392 100644
--- a/lib/axlsx/workbook/worksheet/pivot_table.rb
+++ b/lib/axlsx/workbook/worksheet/pivot_table.rb
@@ -26,17 +26,33 @@ module Axlsx
       @pages = []
       @subtotal = nil
       @no_subtotals_on_headers = []
+      @sort_on_headers = {}
       @style_info = {}
       parse_options options
       yield self if block_given?
     end
 
-    # Defines the headers in which subtotals are not to be included
-    # @return[Array]
+    # Defines the headers in which subtotals are not to be included.
+    # @return [Array]
     attr_accessor :no_subtotals_on_headers
 
+    # Defines the headers in which sort is applied.
+    # Can be an array of headers to sort ascending by default, or a hash for specific control
+    # (with headers as keys, `:ascending` or `:descending` as values).
+    #
+    # Examples: `["year", "month"]` or `{"year" => :descending, "month" => :descending}`
+    # @return [Hash]
+    attr_reader :sort_on_headers
+
+    # (see #sort_on_headers)
+    def sort_on_headers=(headers)
+      headers ||= {}
+      headers = Hash[*headers.map { |h| [h, :ascending] }.flatten] if headers.is_a?(Array)
+      @sort_on_headers = headers
+    end
+
     # Style info for the pivot table
-    # @return[Hash]
+    # @return [Hash]
     attr_accessor :style_info
 
     # The reference to the table data
@@ -178,9 +194,13 @@ module Axlsx
 
       str << ('<location firstDataCol="1" firstDataRow="1" firstHeaderRow="1" ref="' << ref << '"/>')
       str << ('<pivotFields count="' << header_cells_count.to_s << '">')
+
       header_cell_values.each do |cell_value|
-        str << pivot_field_for(cell_value, !no_subtotals_on_headers.include?(cell_value))
+        subtotal = !no_subtotals_on_headers.include?(cell_value)
+        sorttype = sort_on_headers[cell_value]
+        str << pivot_field_for(cell_value, subtotal, sorttype)
       end
+
       str << '</pivotFields>'
       if rows.empty?
         str << '<rowFields count="1"><field x="-2"/></rowFields>'
@@ -267,22 +287,31 @@ module Axlsx
 
     private
 
-    def pivot_field_for(cell_ref, subtotal=true)
+    def pivot_field_for(cell_ref, subtotal, sorttype)
+      attributes = %w[compact="0" outline="0" subtotalTop="0" showAll="0" includeNewItemsInFilter="1"]
+      items_tag = '<items count="1"><item t="default"/></items>'
+      include_items_tag = false
+
       if rows.include? cell_ref
+        attributes << 'axis="axisRow"'
+        attributes << "sortType=\"#{sorttype == :descending ? 'descending' : 'ascending'}\"" if sorttype
         if subtotal
-          '<pivotField axis="axisRow" compact="0" outline="0" subtotalTop="0" showAll="0" includeNewItemsInFilter="1"><items count="1"><item t="default"/></items></pivotField>'
+          include_items_tag = true
         else
-          '<pivotField axis="axisRow" compact="0" outline="0" subtotalTop="0" showAll="0" includeNewItemsInFilter="1" defaultSubtotal="0"></pivotField>'
+          attributes << 'defaultSubtotal="0"'
         end
       elsif columns.include? cell_ref
-        '<pivotField axis="axisCol" compact="0" outline="0" subtotalTop="0" showAll="0" includeNewItemsInFilter="1"><items count="1"><item t="default"/></items></pivotField>'
+        attributes << 'axis="axisCol"'
+        attributes << "sortType=\"#{sorttype == :descending ? 'descending' : 'ascending'}\"" if sorttype
+        include_items_tag = true
       elsif pages.include? cell_ref
-        '<pivotField axis="axisPage" compact="0" outline="0" subtotalTop="0" showAll="0" includeNewItemsInFilter="1"><items count="1"><item t="default"/></items></pivotField>'
+        attributes << 'axis="axisPage"'
+        include_items_tag = true
       elsif data_refs.include? cell_ref
-        '<pivotField dataField="1" compact="0" outline="0" subtotalTop="0" showAll="0" includeNewItemsInFilter="1"></pivotField>'
-      else
-        '<pivotField compact="0" outline="0" subtotalTop="0" showAll="0" includeNewItemsInFilter="1"></pivotField>'
+        attributes << 'dataField="1"'
       end
+
+      "<pivotField #{attributes.join(' ')}>#{include_items_tag ? items_tag : nil}</pivotField>"
     end
 
     def data_refs
diff --git a/test/workbook/worksheet/tc_pivot_table.rb b/test/workbook/worksheet/tc_pivot_table.rb
index 591f0ba9..c7aa53cc 100644
--- a/test/workbook/worksheet/tc_pivot_table.rb
+++ b/test/workbook/worksheet/tc_pivot_table.rb
@@ -87,6 +87,19 @@ class TestPivotTable < Test::Unit::TestCase
     assert_equal(['Year'], pivot_table.no_subtotals_on_headers)
   end
 
+  def test_add_pivot_table_with_months_sorted
+    pivot_table = @ws.add_pivot_table('G5:G6', 'A1:E5', {:sort_on_headers=>['Month']}) do |pt|
+      pt.data = ['Sales']
+      pt.rows = ['Year','Month']
+    end
+    assert_equal({'Month' => :ascending}, pivot_table.sort_on_headers)
+
+    pivot_table.sort_on_headers = {'Month' => :descending}
+    assert_equal({'Month' => :descending}, pivot_table.sort_on_headers)
+
+    shared_test_pivot_table_xml_validity(pivot_table)
+  end
+
   def test_header_indices
     pivot_table = @ws.add_pivot_table('G5:G6', 'A1:E5')
     assert_equal(0,   pivot_table.header_index_of('Year'   ))