Add support for custom table total functions.

Feature request #982
jmcnamara · Sep 7, 2023 · 3320f82 · 3320f82
1 parent ab13807
commit 3320f82
Show file tree

Hide file tree

Showing 10 changed files with 188 additions and 27 deletions.
diff --git a/xlsxwriter/table.py b/xlsxwriter/table.py
@@ -148,11 +148,16 @@ def _write_table_column(self, col_data):
         if "format" in col_data and col_data["format"] is not None:
             attributes.append(("dataDxfId", col_data["format"]))
 
-        if col_data.get("formula"):
+        if col_data.get("formula") or col_data.get("custom_total"):
             self._xml_start_tag("tableColumn", attributes)
 
-            # Write the calculatedColumnFormula element.
-            self._write_calculated_column_formula(col_data["formula"])
+            if col_data.get("formula"):
+                # Write the calculatedColumnFormula element.
+                self._write_calculated_column_formula(col_data["formula"])
+
+            if col_data.get("custom_total"):
+                # Write the totalsRowFormula element.
+                self._write_totals_row_formula(col_data.get("custom_total"))
 
             self._xml_end_tag("tableColumn")
         else:
@@ -182,3 +187,7 @@ def _write_table_style_info(self):
     def _write_calculated_column_formula(self, formula):
         # Write the <calculatedColumnFormula> element.
         self._xml_data_element("calculatedColumnFormula", formula)
+
+    def _write_totals_row_formula(self, formula):
+        # Write the <totalsRowFormula> element.
+        self._xml_data_element("totalsRowFormula", formula)
diff --git a/xlsxwriter/test/comparison/test_table09.py b/xlsxwriter/test/comparison/test_table09.py
@@ -57,13 +57,13 @@ def test_create_file(self):
                 "columns": [
                     {"total_string": "Total"},
                     {},
-                    {"total_function": "Average"},
-                    {"total_function": "COUNT"},
+                    {"total_function": "average"},
+                    {"total_function": "count"},
                     {"total_function": "count_nums"},
                     {"total_function": "max"},
                     {"total_function": "min"},
                     {"total_function": "sum"},
-                    {"total_function": "std Dev"},
+                    {"total_function": "stdDev"},
                     {"total_function": "var"},
                 ],
             },

diff --git a/xlsxwriter/test/comparison/test_table10.py b/xlsxwriter/test/comparison/test_table10.py
@@ -58,13 +58,13 @@ def test_create_file(self):
                 "columns": [
                     {"total_string": "Total"},
                     {},
-                    {"total_function": "Average"},
-                    {"total_function": "COUNT"},
+                    {"total_function": "average"},
+                    {"total_function": "count"},
                     {"total_function": "count_nums"},
                     {"total_function": "max"},
                     {"total_function": "min"},
                     {"total_function": "sum"},
-                    {"total_function": "std Dev"},
+                    {"total_function": "stdDev"},
                     {
                         "total_function": "var",
                         "formula": "SUM(Table1[[#This Row],[Column1]:[Column3]])",

diff --git a/xlsxwriter/test/comparison/test_table17.py b/xlsxwriter/test/comparison/test_table17.py
@@ -61,13 +61,13 @@ def test_create_file(self):
                 "columns": [
                     {"total_string": "Total"},
                     {},
-                    {"total_function": "Average"},
-                    {"total_function": "COUNT"},
+                    {"total_function": "average"},
+                    {"total_function": "count"},
                     {"total_function": "count_nums"},
                     {"total_function": "max", "total_value": 5},
                     {"total_function": "min"},
                     {"total_function": "sum", "total_value": 3},
-                    {"total_function": "std Dev"},
+                    {"total_function": "stdDev"},
                     {"total_function": "var"},
                 ],
             },

diff --git a/xlsxwriter/test/comparison/test_table32.py b/xlsxwriter/test/comparison/test_table32.py
@@ -0,0 +1,58 @@
+###############################################################################
+#
+# Tests for XlsxWriter.
+#
+# SPDX-License-Identifier: BSD-2-Clause
+# Copyright (c), 2013-2023, John McNamara, jmcnamara@cpan.org
+#
+
+from ..excel_comparison_test import ExcelComparisonTest
+from ...workbook import Workbook
+
+
+class TestCompareXLSXFiles(ExcelComparisonTest):
+    """
+    Test file created by XlsxWriter against a file created by Excel.
+
+    """
+
+    def setUp(self):
+        self.set_filename("table32.xlsx")
+
+        self.ignore_files = [
+            "xl/calcChain.xml",
+            "[Content_Types].xml",
+            "xl/_rels/workbook.xml.rels",
+        ]
+
+    def test_create_file(self):
+        """Test the creation of a simple XlsxWriter file with tables."""
+
+        workbook = Workbook(self.got_filename)
+
+        worksheet = workbook.add_worksheet()
+
+        worksheet.set_column("C:F", 10.288)
+
+        worksheet.write_string("A1", "Column1")
+        worksheet.write_string("B1", "Column2")
+        worksheet.write_string("C1", "Column3")
+        worksheet.write_string("D1", "Column4")
+        worksheet.write_string("E1", "Total")
+
+        worksheet.add_table(
+            "C3:F14",
+            {
+                "total_row": 1,
+                "columns": [
+                    {"total_string": "Total"},
+                    {"total_function": "D5+D9"},
+                    {"total_function": "=SUM([Column3])"},
+                    {"total_function": "count"},
+                ],
+            },
+        )
+
+        workbook.close()
+
+        self.assertExcelEqual()
diff --git a/xlsxwriter/test/comparison/test_table33.py b/xlsxwriter/test/comparison/test_table33.py
@@ -0,0 +1,79 @@
+###############################################################################
+#
+# Tests for XlsxWriter.
+#
+# SPDX-License-Identifier: BSD-2-Clause
+# Copyright (c), 2013-2023, John McNamara, jmcnamara@cpan.org
+#
+
+from ..excel_comparison_test import ExcelComparisonTest
+from ...workbook import Workbook
+
+
+class TestCompareXLSXFiles(ExcelComparisonTest):
+    """
+    Test file created by XlsxWriter against a file created by Excel.
+
+    """
+
+    def setUp(self):
+        self.set_filename("table33.xlsx")
+
+        self.ignore_files = [
+            "xl/calcChain.xml",
+            "[Content_Types].xml",
+            "xl/_rels/workbook.xml.rels",
+        ]
+
+    def test_create_file(self):
+        """Test the creation of a simple XlsxWriter file with tables."""
+
+        workbook = Workbook(self.got_filename)
+
+        worksheet = workbook.add_worksheet()
+        xformat = workbook.add_format({"num_format": 2})
+
+        worksheet.set_column("B:K", 10.288)
+
+        worksheet.write_string("A1", "Column1")
+        worksheet.write_string("B1", "Column2")
+        worksheet.write_string("C1", "Column3")
+        worksheet.write_string("D1", "Column4")
+        worksheet.write_string("E1", "Column5")
+        worksheet.write_string("F1", "Column6")
+        worksheet.write_string("G1", "Column7")
+        worksheet.write_string("H1", "Column8")
+        worksheet.write_string("I1", "Column9")
+        worksheet.write_string("J1", "Column10")
+        worksheet.write_string("K1", "Total")
+
+        data = [0, 0, 0, None, None, 0, 0, 0, 0, 0]
+        worksheet.write_row("B4", data)
+        worksheet.write_row("B5", data)
+
+        worksheet.add_table(
+            "B3:K6",
+            {
+                "total_row": 1,
+                "columns": [
+                    {"total_string": "Total"},
+                    {},
+                    {"total_function": "average"},
+                    {"total_function": "count"},
+                    {"total_function": "count_nums"},
+                    {"total_function": "max"},
+                    {"total_function": "min"},
+                    {"total_function": "sum"},
+                    {"total_function": "std_dev"},
+                    {
+                        "total_function": "=SUM([Column10])",
+                        "formula": "SUM(Table1[[#This Row],[Column1]:[Column3]])",
+                        "format": xformat,
+                    },
+                ],
+            },
+        )
+
+        workbook.close()
+
+        self.assertExcelEqual()
diff --git a/xlsxwriter/test/comparison/xlsx_files/table32.xlsx b/xlsxwriter/test/comparison/xlsx_files/table32.xlsx
diff --git a/xlsxwriter/test/comparison/xlsx_files/table33.xlsx b/xlsxwriter/test/comparison/xlsx_files/table33.xlsx
diff --git a/xlsxwriter/test/table/test_table09.py b/xlsxwriter/test/table/test_table09.py
@@ -37,13 +37,13 @@ def test_assemble_xml_file(self):
                 "columns": [
                     {"total_string": "Total"},
                     {},
-                    {"total_function": "Average"},
-                    {"total_function": "COUNT"},
+                    {"total_function": "average"},
+                    {"total_function": "count"},
                     {"total_function": "count_nums"},
                     {"total_function": "max"},
                     {"total_function": "min"},
                     {"total_function": "sum"},
-                    {"total_function": "std Dev"},
+                    {"total_function": "stdDev"},
                     {"total_function": "var"},
                 ],
             },

diff --git a/xlsxwriter/worksheet.py b/xlsxwriter/worksheet.py
@@ -3342,6 +3342,7 @@ def add_table(self, first_row, first_col, last_row, last_col, options=None):
                 "name": "Column" + str(col_id),
                 "total_string": "",
                 "total_function": "",
+                "custom_total": "",
                 "total_value": 0,
                 "formula": "",
                 "format": None,
@@ -3394,23 +3395,37 @@ def add_table(self, first_row, first_col, last_row, last_col, options=None):
                     # Handle the function for the total row.
                     if user_data.get("total_function"):
                         function = user_data["total_function"]
-
-                        # Massage the function name.
-                        function = function.lower()
-                        function = function.replace("_", "")
-                        function = function.replace(" ", "")
-
-                        if function == "countnums":
+                        if function == "count_nums":
                             function = "countNums"
-                        if function == "stddev":
+                        if function == "std_dev":
                             function = "stdDev"
 
-                        col_data["total_function"] = function
-
-                        formula = self._table_function_to_formula(
-                            function, col_data["name"]
+                        subtotals = set(
+                            [
+                                "average",
+                                "countNums",
+                                "count",
+                                "max",
+                                "min",
+                                "stdDev",
+                                "sum",
+                                "var",
+                            ]
                         )
 
+                        if function in subtotals:
+                            formula = self._table_function_to_formula(
+                                function, col_data["name"]
+                            )
+                        else:
+                            formula = function
+                            if formula.startswith("="):
+                                formula = formula.lstrip("=")
+                            col_data["custom_total"] = formula
+                            function = "custom"
+
+                        col_data["total_function"] = function
+
                         value = user_data.get("total_value", 0)
 
                         self._write_formula(last_row, col_num, formula, xformat, value)