From a9bbdc01826f9f4a17c240a807e7735ec593efdf Mon Sep 17 00:00:00 2001
From: Shivam Mishra <scm.mymail@gmail.com>
Date: Tue, 18 Aug 2020 13:42:18 +0530
Subject: [PATCH] fix: handle both list and dict results

---
 frappe/desk/query_report.py | 37 ++++++++++++++++++++++++++++++++-----
 1 file changed, 32 insertions(+), 5 deletions(-)

diff --git a/frappe/desk/query_report.py b/frappe/desk/query_report.py
index d0a32ef076..656b49046a 100644
--- a/frappe/desk/query_report.py
+++ b/frappe/desk/query_report.py
@@ -216,20 +216,32 @@ def add_data_to_custom_columns(columns, result):
 	return data
 
 def reorder_data_for_custom_columns(custom_columns, columns, result, report_type):
-	custom_column_labels = [col["label"] for col in custom_columns]
+	if not result:
+		return []
 
 	if report_type == 'Query Report':
+		# Assume list result for query reports
+		# Query report columns exclisvuely use Label
+		custom_column_labels = [col["label"] for col in custom_columns]
 		original_column_labels = [col.split(":")[0] for col in columns]
-	else:
-		original_column_labels = [col["label"] for col in columns]
+		return get_columns_from_list(custom_column_labels, original_column_labels, result)
 
+	if type(result[0]) == "list":
+		# If the result is a list of lists
+		original_column_names = [col["fieldname"] for col in columns]
+		return get_columns_from_list(custom_column_names, original_column_names, result)
+	else:
+		# If the result is a list of dicts
+		return get_columns_from_dict(custom_column_names, result)
+
+def get_columns_from_list(columns, target_columns, result):
 	reordered_result = []
 
 	for res in result:
 		r = []
-		for col_name in custom_column_labels:
+		for col_name in columns:
 			try:
-				idx = original_column_labels.index(col_name)
+				idx = target_columns.index(col_name)
 				r.append(res[idx])
 			except ValueError:
 				pass
@@ -238,6 +250,21 @@ def reorder_data_for_custom_columns(custom_columns, columns, result, report_type
 
 	return reordered_result
 
+def get_columns_from_dict(columns, result):
+	reordered_result = []
+
+	for res in result:
+		r = {}
+		for col_name in columns:
+			try:
+				r[col_name] = res[col_name]
+			except KeyError:
+				pass
+
+		reordered_result.append(r)
+
+	return reordered_result
+
 def get_prepared_report_result(report, filters, dn="", user=None):
 	latest_report_data = {}
 	doc = None