feat: improve bench execute argument passing (#35964)

* fix: better way to pass arguments to bench execute

* refactor: move execute logic in separate file

---------

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

frappe/commands/execute.py

@@ -0,0 +1,103 @@
+import json
+
+import frappe
+from frappe.exceptions import SiteNotSpecifiedError
+from frappe.utils.bench_helper import CliCtxObj
+
+
+def _execute(context: CliCtxObj, method, args=None, kwargs=None, profile=False, extra_args=None):
+	for site in context.sites:
+		ret = ""
+		try:
+			frappe.init(site)
+			frappe.connect()
+
+			if args:
+				try:
+					fn_args = eval(args)
+				except NameError:
+					fn_args = [args]
+			else:
+				fn_args = ()
+
+			if kwargs:
+				fn_kwargs = eval(kwargs)
+			else:
+				fn_kwargs = {}
+
+			if extra_args:
+				# parse extra_args
+				# if it starts with --, it is a kwarg
+				# otherwise it is an arg
+				# if it is a kwarg, the next argument is the value
+				# if the next argument starts with --, the value is True
+				# if there is no next argument, the value is True
+
+				# examples:
+				# bench execute method arg1 arg2 -> args=[arg1, arg2]
+				# bench execute method --a 1 --b 2 -> kwargs={a: 1, b: 2}
+				# bench execute method arg1 --a 1 -> args=[arg1], kwargs={a: 1}
+
+				# we need to convert values to python objects if possible
+				def parse_value(value):
+					try:
+						return json.loads(value)
+					except Exception:
+						return value
+
+				extra_args = list(extra_args)
+				while extra_args:
+					arg = extra_args.pop(0)
+					if arg.startswith("--"):
+						key = arg[2:]
+						if extra_args and not extra_args[0].startswith("--"):
+							value = parse_value(extra_args.pop(0))
+						else:
+							value = True
+						fn_kwargs[key] = value
+					else:
+						fn_args += (parse_value(arg),)
+
+			pr = None
+			if profile:
+				import cProfile
+
+				pr = cProfile.Profile()
+				pr.enable()
+
+			try:
+				fn = frappe.get_attr(method)
+			except Exception:
+				fn = None
+
+			if fn:
+				ret = fn(*fn_args, **fn_kwargs)
+			else:
+				# eval is safe here because input is from console
+				code = compile(method, "<bench execute>", "eval")
+				ret = eval(code, globals(), locals())  # nosemgrep
+				if callable(ret):
+					suffix = "(*fn_args, **fn_kwargs)"
+					code = compile(method + suffix, "<bench execute>", "eval")
+					ret = eval(code, globals(), locals())  # nosemgrep
+
+			if profile and pr:
+				import pstats
+				from io import StringIO
+
+				pr.disable()
+				s = StringIO()
+				pstats.Stats(pr, stream=s).sort_stats("cumulative").print_stats(0.5)
+				print(s.getvalue())
+
+			if frappe.db:
+				frappe.db.commit()
+		finally:
+			frappe.destroy()
+		if ret:
+			from frappe.utils.response import json_handler
+
+			print(json.dumps(ret, default=json_handler).strip('"'))
+
+	if not context.sites:
+		raise SiteNotSpecifiedError

frappe/commands/test_commands.py

@@ -244,6 +244,26 @@ class TestCommands(BaseTestCommands):
 		self.assertEqual(self.returncode, 0)
 		self.assertEqual(self.stdout, frappe.bold(text="DocType"))
 
+		# test 5: execute a command with extra args
+		self.execute("bench --site {site} execute frappe.bold DocType")
+		self.assertEqual(self.returncode, 0)
+		self.assertEqual(self.stdout, frappe.bold(text="DocType"))
+
+		# test 6: execute a command with extra kwargs
+		self.execute("bench --site {site} execute frappe.bold --text DocType")
+		self.assertEqual(self.returncode, 0)
+		self.assertEqual(self.stdout, frappe.bold(text="DocType"))
+
+		# test 7: execute a command with extra args and kwargs
+		self.execute("bench --site {site} execute frappe.utils.add_to_date '2024-01-01' --days 1")
+		self.assertEqual(self.returncode, 0)
+		self.assertEqual(self.stdout, "2024-01-02")
+
+		# test 8: execute a command with extra args and kwargs with types
+		self.execute("bench --site {site} execute frappe.utils.add_to_date --date '2024-01-01' --days 1")
+		self.assertEqual(self.returncode, 0)
+		self.assertEqual(self.stdout, "2024-01-02")
+
 	@skipIf(
 		frappe.conf.db_type == "sqlite",
 		"Not for SQLite for now",

frappe/commands/utils.py

@@ -247,70 +247,18 @@ def reset_perms(context: CliCtxObj):
 		raise SiteNotSpecifiedError
 
 
-@click.command("execute")
+@click.command("execute", context_settings=EXTRA_ARGS_CTX)
 @click.argument("method")
 @click.option("--args")
 @click.option("--kwargs")
 @click.option("--profile", is_flag=True, default=False)
+@click.argument("extra_args", nargs=-1)
 @pass_context
-def execute(context: CliCtxObj, method, args=None, kwargs=None, profile=False):
+def execute(context: CliCtxObj, method, args=None, kwargs=None, profile=False, extra_args=None):
 	"Execute a function"
-	for site in context.sites:
-		ret = ""
-		try:
-			frappe.init(site)
-			frappe.connect()
+	from frappe.commands.execute import _execute
 
-			if args:
-				try:
-					fn_args = eval(args)
-				except NameError:
-					fn_args = [args]
-			else:
-				fn_args = ()
-
-			if kwargs:
-				fn_kwargs = eval(kwargs)
-			else:
-				fn_kwargs = {}
-
-			if profile:
-				import cProfile
-
-				pr = cProfile.Profile()
-				pr.enable()
-
-			try:
-				ret = frappe.get_attr(method)(*fn_args, **fn_kwargs)
-			except Exception:
-				# eval is safe here because input is from console
-				code = compile(method, "<bench execute>", "eval")
-				ret = eval(code, globals(), locals())  # nosemgrep
-				if callable(ret):
-					suffix = "(*fn_args, **fn_kwargs)"
-					code = compile(method + suffix, "<bench execute>", "eval")
-					ret = eval(code, globals(), locals())  # nosemgrep
-
-			if profile:
-				import pstats
-				from io import StringIO
-
-				pr.disable()
-				s = StringIO()
-				pstats.Stats(pr, stream=s).sort_stats("cumulative").print_stats(0.5)
-				print(s.getvalue())
-
-			if frappe.db:
-				frappe.db.commit()
-		finally:
-			frappe.destroy()
-		if ret:
-			from frappe.utils.response import json_handler
-
-			print(json.dumps(ret, default=json_handler).strip('"'))
-
-	if not context.sites:
-		raise SiteNotSpecifiedError
+	_execute(context, method, args, kwargs, profile, extra_args)
 
 
 @click.command("add-to-email-queue")