fix if custom_cols is None

fixes bug introduced in #40
wharton · Nov 9, 2021 · 82af43b · 82af43b
1 parent cce21fb
commit 82af43b
Showing 1 changed file with 10 additions and 8 deletions.
diff --git a/drf_renderer_xlsx/renderers.py b/drf_renderer_xlsx/renderers.py
@@ -163,14 +163,16 @@ def render(self, data, accepted_media_type=None, renderer_context=None):
             self.xlsx_header_dict = self._flatten_serializer_keys(
                 drf_view.get_serializer(), use_labels=use_labels
             )
-
-            custom_header_dict = {
-                key: self.custom_cols[key].get('label', None) or key 
-                for key in self.custom_cols.keys()
-            }
-            self.combined_header_dict = dict(
-                list(self.xlsx_header_dict.items()) + list(custom_header_dict.items())
-            )
+            if self.custom_cols:
+                custom_header_dict = {
+                    key: self.custom_cols[key].get('label', None) or key
+                    for key in self.custom_cols.keys()
+                }
+                self.combined_header_dict = dict(
+                    list(self.xlsx_header_dict.items()) + list(custom_header_dict.items())
+                )
+            else:
+                self.combined_header_dict = self.xlsx_header_dict
 
             for column_name, column_label in self.combined_header_dict.items():
                 if column_name == "row_color":