From 848bad2ce8aff6c6ff6c54eda3dc3437439b1991 Mon Sep 17 00:00:00 2001
From: antazoey <jules@apeworx.io>
Date: Fri, 3 Nov 2023 12:13:42 -0500
Subject: [PATCH] feat: allow argument overrides (#1723)

---
 src/ape/cli/arguments.py     | 14 +++++++---
 tests/functional/test_cli.py | 52 ++++++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 4 deletions(-)

diff --git a/src/ape/cli/arguments.py b/src/ape/cli/arguments.py
index 8e8c671ba5..eba1a22b47 100644
--- a/src/ape/cli/arguments.py
+++ b/src/ape/cli/arguments.py
@@ -24,24 +24,30 @@ def _alias_callback(ctx, param, value):
     return value
 
 
-def existing_alias_argument(account_type: _ACCOUNT_TYPE_FILTER = None):
+def existing_alias_argument(account_type: _ACCOUNT_TYPE_FILTER = None, **kwargs):
     """
     A ``click.argument`` for an existing account alias.
 
     Args:
         account_type (Type[:class:`~ape.api.accounts.AccountAPI`], optional):
           If given, limits the type of account the user may choose from.
+        **kwargs: click.argument overrides.
     """
 
-    return click.argument("alias", type=Alias(account_type=account_type))
+    type_ = kwargs.pop("type", Alias(account_type=account_type))
+    return click.argument("alias", type=type_, **kwargs)
 
 
-def non_existing_alias_argument():
+def non_existing_alias_argument(**kwargs):
     """
     A ``click.argument`` for an account alias that does not yet exist in ape.
+
+    Args:
+        **kwargs: click.argument overrides.
     """
 
-    return click.argument("alias", callback=_alias_callback)
+    callback = kwargs.pop("callback", _alias_callback)
+    return click.argument("alias", callback=callback, **kwargs)
 
 
 def _create_contracts_paths(ctx, param, value):
diff --git a/tests/functional/test_cli.py b/tests/functional/test_cli.py
index ff1fbc00e0..3437760f4e 100644
--- a/tests/functional/test_cli.py
+++ b/tests/functional/test_cli.py
@@ -9,8 +9,10 @@
     PromptChoice,
     account_option,
     contract_file_paths_argument,
+    existing_alias_argument,
     get_user_selected_account,
     network_option,
+    non_existing_alias_argument,
     verbosity_option,
 )
 from ape.exceptions import AccountsError
@@ -346,3 +348,53 @@ def cmd(file_paths):
 
     result = runner.invoke(cmd, ["path0", "path1"])
     assert "Contract 'path0' not found" in result.output
+
+
+def test_existing_alias_option(runner):
+    @click.command()
+    @existing_alias_argument()
+    def cmd(alias):
+        click.echo(alias)
+
+    result = runner.invoke(cmd, ["TEST::0"])
+    assert "TEST::0" in result.output
+
+
+def test_existing_alias_option_custom_callback(runner):
+    magic_value = "THIS IS A TEST"
+
+    def custom_callback(*args, **kwargs):
+        return magic_value
+
+    @click.command()
+    @existing_alias_argument(callback=custom_callback)
+    def cmd(alias):
+        click.echo(alias)
+
+    result = runner.invoke(cmd, ["TEST::0"])
+    assert magic_value in result.output
+
+
+def test_non_existing_alias_option(runner):
+    @click.command()
+    @non_existing_alias_argument()
+    def cmd(alias):
+        click.echo(alias)
+
+    result = runner.invoke(cmd, ["non-exists"])
+    assert "non-exists" in result.output
+
+
+def test_non_existing_alias_option_custom_callback(runner):
+    magic_value = "THIS IS A TEST"
+
+    def custom_callback(*args, **kwargs):
+        return magic_value
+
+    @click.command()
+    @non_existing_alias_argument(callback=custom_callback)
+    def cmd(alias):
+        click.echo(alias)
+
+    result = runner.invoke(cmd, ["non-exists"])
+    assert magic_value in result.output