refactor: patch for migration

erpnext/patches/v16_0/migrate_budget_records_to_new_structure.py

@@ -3,61 +3,93 @@ from frappe.utils import add_months, flt, get_first_day, get_last_day
 
 
 def execute():
-	budgets = frappe.get_all("Budget", filters={"docstatus": ["in", [0, 1]]}, pluck="name")
+	remove_old_property_setter()
 
-	for budget in budgets:
+	budget_names = frappe.db.get_list(
-		old_budget = frappe.get_doc("Budget", budget)
+		"Budget",
+		filters={"docstatus": ["in", [0, 1]]},
+		pluck="name",
+	)
 
-		accounts = frappe.get_all(
+	for budget in budget_names:
+		migrate_single_budget(budget)
+
+
+def remove_old_property_setter():
+	old_property_setter = frappe.db.get_value(
+		"Property Setter",
+		{
+			"doc_type": "Budget",
+			"field_name": "naming_series",
+			"property": "options",
+			"value": "Budget-.YYYY.-",
+		},
+		"name",
+	)
+
+	if old_property_setter:
+		frappe.delete_doc("Property Setter", old_property_setter, force=1)
+
+
+def migrate_single_budget(budget_name):
+	budget_doc = frappe.get_doc("Budget", budget_name)
+
+	account_rows = frappe.get_all(
 		"Budget Account",
-			filters={"parent": old_budget.name},
+		filters={"parent": budget_name},
 		fields=["account", "budget_amount"],
 		order_by="idx asc",
 	)
 
-		if not accounts:
+	if not account_rows:
-			continue
+		return
 
-		old_distribution = []
+	frappe.db.delete("Budget Account", {"parent": budget_doc.name})
-		if old_budget.monthly_distribution:
+
-			old_distribution = frappe.get_all(
+	percentage_allocations = get_percentage_allocations(budget_doc)
-				"Monthly Distribution Percentage",
+
-				filters={"parent": old_budget.monthly_distribution},
+	fiscal_year = frappe.get_cached_value(
-				fields=["percentage_allocation"],
+		"Fiscal Year",
-				order_by="idx asc",
+		budget_doc.fiscal_year,
+		["name", "year_start_date", "year_end_date"],
+		as_dict=True,
 	)
 
-		if old_distribution:
+	for row in account_rows:
-			percentage_list = [flt(d.percentage_allocation) for d in old_distribution]
+		create_new_budget_from_row(budget_doc, fiscal_year, row, percentage_allocations)
+
+	if budget_doc.docstatus == 1:
+		budget_doc.cancel()
 	else:
-			percentage_list = [100 / 12] * 12
+		frappe.delete_doc("Budget", budget_name)
 
-		fy = frappe.get_doc("Fiscal Year", old_budget.fiscal_year)
-		fy_start = fy.year_start_date
-		fy_end = fy.year_end_date
 
-		for account in accounts:
+def get_percentage_allocations(budget_doc):
-			new = frappe.new_doc("Budget")
+	if budget_doc.monthly_distribution:
+		distribution_doc = frappe.get_cached_doc("Monthly Distribution", budget_doc.monthly_distribution)
+		return [flt(row.percentage_allocation) for row in distribution_doc.percentages]
 
-			new.naming_series = "BUDGET-.########"
+	return [100 / 12] * 12
-			new.budget_against = old_budget.budget_against
-			new.company = old_budget.company
-			new.cost_center = old_budget.cost_center
-			new.project = old_budget.project
-			new.fiscal_year = fy.name
 
-			new.from_fiscal_year = fy.name
-			new.to_fiscal_year = fy.name
-			new.budget_start_date = fy_start
-			new.budget_end_date = fy_end
 
-			new.account = account.account
+def create_new_budget_from_row(budget_doc, fiscal_year, account_row, percentage_allocations):
-			new.budget_amount = flt(account.budget_amount)
+	new_budget = frappe.new_doc("Budget")
-			new.distribution_frequency = "Monthly"
 
-			new.distribute_equally = 1 if len(set(percentage_list)) == 1 else 0
+	core_fields = ["budget_against", "company", "cost_center", "project"]
+	for field in core_fields:
+		new_budget.set(field, budget_doc.get(field))
 
-			fields_to_copy = [
+	new_budget.from_fiscal_year = fiscal_year.name
+	new_budget.to_fiscal_year = fiscal_year.name
+	new_budget.budget_start_date = fiscal_year.year_start_date
+	new_budget.budget_end_date = fiscal_year.year_end_date
+
+	new_budget.account = account_row.account
+	new_budget.budget_amount = flt(account_row.budget_amount)
+	new_budget.distribution_frequency = "Monthly"
+	new_budget.distribute_equally = 1 if len(set(percentage_allocations)) == 1 else 0
+
+	copy_fields = [
 		"applicable_on_material_request",
 		"action_if_annual_budget_exceeded_on_mr",
 		"action_if_accumulated_monthly_budget_exceeded_on_mr",
@@ -72,36 +104,25 @@ def execute():
 		"action_if_accumulated_monthly_exceeded_on_cumulative_expense",
 	]
 
-			for field in fields_to_copy:
+	for field in copy_fields:
-				if hasattr(old_budget, field):
+		new_budget.set(field, budget_doc.get(field))
-					new.set(field, old_budget.get(field))
 
-			start = fy_start
+	current_start = fiscal_year.year_start_date
-			for percentage in percentage_list:
+	for percentage in percentage_allocations:
-				row_start = get_first_day(start)
+		new_budget.append(
-				row_end = get_last_day(start)
-
-				new.append(
 			"budget_distribution",
 			{
-						"start_date": row_start,
+				"start_date": get_first_day(current_start),
-						"end_date": row_end,
+				"end_date": get_last_day(current_start),
 				"percent": percentage,
-						"amount": new.budget_amount * percentage / 100,
+				"amount": new_budget.budget_amount * percentage / 100,
 			},
 		)
+		current_start = add_months(current_start, 1)
 
-				start = add_months(start, 1)
+	new_budget.flags.ignore_validate = True
+	new_budget.flags.ignore_links = True
+	new_budget.insert(ignore_permissions=True, ignore_mandatory=True)
 
-			new.flags.ignore_validate = True
+	if budget_doc.docstatus == 1:
-			new.flags.ignore_links = True
+		new_budget.submit()
-
-			new.insert(ignore_permissions=True, ignore_mandatory=True)
-
-			if old_budget.docstatus == 1:
-				new.submit()
-
-		if old_budget.docstatus == 1:
-			old_budget.cancel()
-		else:
-			old_budget.delete()