refactor: make code more pythonic

erpnext/controllers/budget_controller.py

@@ -1,7 +1,8 @@
+from collections import OrderedDict
+
 import frappe
 from frappe import qb
 
-# from erpnext.accounts.doctype.accounting_dimension.accounting_dimension import get_dimensions
 from erpnext.accounts.utils import get_fiscal_year
 
 
@@ -28,7 +29,7 @@ class BudgetValidation:
 			]
 		)
 
-	def get_budget_records(self):
+	def get_budget_records(self) -> list:
 		bud = qb.DocType("Budget")
 		bud_acc = qb.DocType("Budget Account")
 		query = (
@@ -58,69 +59,59 @@ class BudgetValidation:
 		for x in self.dimensions:
 			query = query.select(bud[x.get("fieldname")])
 
-		self.budgets = query.run(as_dict=True)
+		_budgets = query.run(as_dict=True)
+		return _budgets
 
-	def generate_active_budget_keys(self):
+	def build_budget_keys_and_map(self):
 		"""
 		key structure - (dimension_type, dimension, GL account)
 		"""
-		self.get_budget_records()
+		_budgets = self.get_budget_records()
 		_keys = []
-		for x in self.budgets:
-			budget_against = frappe.scrub(x.budget_against)
-			dimension = x.get(budget_against)
-			_keys.append((budget_against, dimension, x.account))
-		self.active_keys = set(_keys)
+		self.budget_map = OrderedDict()
+		for _bud in _budgets:
+			budget_against = frappe.scrub(_bud.budget_against)
+			dimension = _bud.get(budget_against)
+			key = (budget_against, dimension, _bud.account)
+			# TODO: ensure duplicate keys are not possible
+			self.budget_map[key] = _bud
+		self.budget_keys = self.budget_map.keys()
 
-	def generate_doc_dimension_keys(self):
+	def build_doc_or_item_keys_and_map(self):
 		"""
 		key structure - (dimension_type, dimension, GL account)
 		"""
-		keys = []
-		for itm in self.doc.items:
-			keys.extend(
-				[
-					(dim.get("fieldname"), itm.get(dim.get("fieldname")), itm.expense_account)
-					for dim in self.dimensions
-					if itm.get(dim.get("fieldname"))
-				]
-			)
-		self.item_dimension_keys = set(keys)
-
-	def build_processing_dictionary(self):
-		self.budget_map = frappe._dict()
-
-		for x in self.budgets:
-			budget_against = frappe.scrub(x.budget_against)
-			dimension = x.get(budget_against)
-			key = (budget_against, dimension, x.account)
-			if key in self.overlap:
-				self.budget_map[key] = frappe._dict(
-					{"budget_amount": x.budget_amount, "items_to_process": []}
-				)
-
+		self.doc_or_item_map = OrderedDict()
+		_key = []
 		for itm in self.doc.items:
 			for dim in self.dimensions:
 				if itm.get(dim.get("fieldname")):
 					key = (dim.get("fieldname"), itm.get(dim.get("fieldname")), itm.expense_account)
+					# TODO: How to handle duplicate items - same item with same dimension with same account
+					self.doc_or_item_map.setdefault(key, []).append(itm)
+		self.doc_or_item_keys = self.doc_or_item_map.keys()
 
-					if key in self.overlap:
-						self.budget_map[key]["items_to_process"].append(itm)
+	def build_to_validate_map(self):
+		self.overlap = self.budget_keys & self.doc_or_item_keys
+		self.to_validate = OrderedDict()
+
+		for key in self.overlap:
+			self.to_validate[key] = OrderedDict(
+				{
+					"budget_amount": self.budget_map[key].budget_amount,
+					"items_to_process": self.doc_or_item_map[key],
+				}
+			)
 
 	def validate(self):
-		self.generate_active_budget_keys()
-		self.generate_doc_dimension_keys()
-
-		self.overlap = self.active_keys & self.item_dimension_keys
-		self.build_processing_dictionary()
+		self.build_budget_keys_and_map()
+		self.build_doc_or_item_keys_and_map()
+		self.build_to_validate_map()
 		self.validate_for_overbooking()
 
-	def get_booked_amount(self):
-		pass
-
 	def validate_for_overbooking(self):
-		# Need to fetch historical amount and add them to the current document
-		for k, v in self.budget_map.items():
-			current_amount = sum([x.amount for x in v.items_to_process])
-			self.budget_map[k]["current_amount"] = current_amount
+		# TODO: Need to fetch historical amount and add them to the current document
+		# TODO: handle applicable checkboxes
+		for k, v in self.to_validate.items():
+			current_amount = sum([x.amount for x in v.get("items_to_process")])
 			print((k, v.get("budget_amount"), current_amount))