Adds overwrite flag and updates tests for btwallet3

bittensor_cli/cli.py

@@ -132,15 +132,10 @@ class Options:
     ss58_address = typer.Option(
         None, "--ss58", "--ss58-address", help="The SS58 address of the coldkey."
     )
-    overwrite_coldkey = typer.Option(
+    overwrite = typer.Option(
         False,
-        help="Overwrite the old coldkey with the newly generated coldkey.",
-        prompt=True,
-    )
-    overwrite_hotkey = typer.Option(
-        False,
-        help="Overwrite the old hotkey with the newly generated hotkey.",
-        prompt=True,
+        "--overwrite/--no-overwrite",
+        help="Overwrite the existing wallet file with the new one.",
     )
     network = typer.Option(
         None,
@@ -1725,6 +1720,7 @@ def wallet_regen_coldkey(
         json: Optional[str] = Options.json,
         json_password: Optional[str] = Options.json_password,
         use_password: Optional[bool] = Options.use_password,
+        overwrite: bool = Options.overwrite,
         quiet: bool = Options.quiet,
         verbose: bool = Options.verbose,
     ):
@@ -1770,6 +1766,7 @@ def wallet_regen_coldkey(
                 json,
                 json_password,
                 use_password,
+                overwrite,
             )
         )
 
@@ -1780,6 +1777,7 @@ def wallet_regen_coldkey_pub(
         wallet_hotkey: Optional[str] = Options.wallet_hotkey,
         public_key_hex: Optional[str] = Options.public_hex_key,
         ss58_address: Optional[str] = Options.ss58_address,
+        overwrite: bool = Options.overwrite,
         quiet: bool = Options.quiet,
         verbose: bool = Options.verbose,
     ):
@@ -1826,7 +1824,7 @@ def wallet_regen_coldkey_pub(
             rich.print("[red]Error: Invalid SS58 address or public key![/red]")
             raise typer.Exit()
         return self._run_command(
-            wallets.regen_coldkey_pub(wallet, ss58_address, public_key_hex)
+            wallets.regen_coldkey_pub(wallet, ss58_address, public_key_hex, overwrite)
         )
 
     def wallet_regen_hotkey(
@@ -1842,6 +1840,7 @@ def wallet_regen_hotkey(
             False,  # Overriden to False
             help="Set to 'True' to protect the generated Bittensor key with a password.",
         ),
+        overwrite: bool = Options.overwrite,
         quiet: bool = Options.quiet,
         verbose: bool = Options.verbose,
     ):
@@ -1880,6 +1879,7 @@ def wallet_regen_hotkey(
                 json,
                 json_password,
                 use_password,
+                overwrite,
             )
         )
 
@@ -1898,6 +1898,7 @@ def wallet_new_hotkey(
             False,  # Overriden to False
             help="Set to 'True' to protect the generated Bittensor key with a password.",
         ),
+        overwrite: bool = Options.overwrite,
         quiet: bool = Options.quiet,
         verbose: bool = Options.verbose,
     ):
@@ -1935,7 +1936,9 @@ def wallet_new_hotkey(
             validate=WV.WALLET,
         )
         n_words = get_n_words(n_words)
-        return self._run_command(wallets.new_hotkey(wallet, n_words, use_password))
+        return self._run_command(
+            wallets.new_hotkey(wallet, n_words, use_password, overwrite)
+        )
 
     def wallet_new_coldkey(
         self,
@@ -1949,6 +1952,7 @@ def wallet_new_coldkey(
             help="The number of words used in the mnemonic. Options: [12, 15, 18, 21, 24]",
         ),
         use_password: Optional[bool] = Options.use_password,
+        overwrite: bool = Options.overwrite,
         quiet: bool = Options.quiet,
         verbose: bool = Options.verbose,
     ):
@@ -1985,7 +1989,9 @@ def wallet_new_coldkey(
             validate=WV.NONE,
         )
         n_words = get_n_words(n_words)
-        return self._run_command(wallets.new_coldkey(wallet, n_words, use_password))
+        return self._run_command(
+            wallets.new_coldkey(wallet, n_words, use_password, overwrite)
+        )
 
     def wallet_check_ck_swap(
         self,
@@ -2019,6 +2025,7 @@ def wallet_create_wallet(
         wallet_hotkey: Optional[str] = Options.wallet_hotkey,
         n_words: Optional[int] = None,
         use_password: bool = Options.use_password,
+        overwrite: bool = Options.overwrite,
         quiet: bool = Options.quiet,
         verbose: bool = Options.verbose,
     ):
@@ -2064,6 +2071,7 @@ def wallet_create_wallet(
                 wallet,
                 n_words,
                 use_password,
+                overwrite,
             )
         )
 

bittensor_cli/src/commands/wallets.py

@@ -69,6 +69,7 @@ async def regen_coldkey(
     json_path: Optional[str] = None,
     json_password: Optional[str] = "",
     use_password: Optional[bool] = True,
+    overwrite: Optional[bool] = False,
 ):
     """Creates a new coldkey under this wallet"""
     json_str: Optional[str] = None
@@ -83,7 +84,7 @@ async def regen_coldkey(
             seed=seed,
             json=(json_str, json_password) if all([json_str, json_password]) else None,
             use_password=use_password,
-            overwrite=False,
+            overwrite=overwrite,
         )
 
         if isinstance(new_wallet, Wallet):
@@ -101,13 +102,14 @@ async def regen_coldkey_pub(
     wallet: Wallet,
     ss58_address: str,
     public_key_hex: str,
+    overwrite: Optional[bool] = False,
 ):
     """Creates a new coldkeypub under this wallet."""
     try:
         new_coldkeypub = wallet.regenerate_coldkeypub(
             ss58_address=ss58_address,
             public_key=public_key_hex,
-            overwrite=False,
+            overwrite=overwrite,
         )
         if isinstance(new_coldkeypub, Wallet):
             console.print(
@@ -125,6 +127,7 @@ async def regen_hotkey(
     json_path: Optional[str],
     json_password: Optional[str] = "",
     use_password: Optional[bool] = False,
+    overwrite: Optional[bool] = False,
 ):
     """Creates a new hotkey under this wallet."""
     json_str: Optional[str] = None
@@ -141,7 +144,7 @@ async def regen_hotkey(
             seed=seed,
             json=(json_str, json_password) if all([json_str, json_password]) else None,
             use_password=use_password,
-            overwrite=False,
+            overwrite=overwrite,
         )
         if isinstance(new_hotkey, Wallet):
             console.print(
@@ -158,13 +161,14 @@ async def new_hotkey(
     wallet: Wallet,
     n_words: int,
     use_password: bool,
+    overwrite: Optional[bool] = False,
 ):
     """Creates a new hotkey under this wallet."""
     try:
         wallet.create_new_hotkey(
             n_words=n_words,
             use_password=use_password,
-            overwrite=False,
+            overwrite=overwrite,
         )
     except KeyFileError:
         print_error("KeyFileError: File is not writable")
@@ -174,13 +178,14 @@ async def new_coldkey(
     wallet: Wallet,
     n_words: int,
     use_password: bool,
+    overwrite: Optional[bool] = False,
 ):
     """Creates a new coldkey under this wallet."""
     try:
         wallet.create_new_coldkey(
             n_words=n_words,
             use_password=use_password,
-            overwrite=False,
+            overwrite=overwrite,
         )
     except KeyFileError:
         print_error("KeyFileError: File is not writable")
@@ -190,13 +195,14 @@ async def wallet_create(
     wallet: Wallet,
     n_words: int = 12,
     use_password: bool = True,
+    overwrite: Optional[bool] = False,
 ):
     """Creates a new wallet."""
     try:
         wallet.create_new_coldkey(
             n_words=n_words,
             use_password=use_password,
-            overwrite=False,
+            overwrite=overwrite,
         )
     except KeyFileError:
         print_error("KeyFileError: File is not writable")
@@ -205,7 +211,7 @@ async def wallet_create(
         wallet.create_new_hotkey(
             n_words=n_words,
             use_password=False,
-            overwrite=False,
+            overwrite=overwrite,
         )
     except KeyFileError:
         print_error("KeyFileError: File is not writable")

requirements.txt

@@ -16,5 +16,5 @@ scalecodec==1.2.11
 substrate-interface~=1.7.9
 typer~=0.12
 websockets>=14.1
-bittensor-wallet>=2.1.3
+bittensor-wallet>=3.0.0
 bt-decode==0.4.0

tests/e2e_tests/test_wallet_creations.py

@@ -304,7 +304,7 @@ def test_wallet_creations(wallet_setup):
     assert wallet_status, message
 
 
-def test_wallet_regen(wallet_setup):
+def test_wallet_regen(wallet_setup, capfd):
     """
     Test the regeneration of coldkeys, hotkeys, and coldkeypub files using mnemonics or ss58 address.
 
@@ -342,7 +342,8 @@ def test_wallet_regen(wallet_setup):
     # Verify the command has output, as expected
     assert result.stdout is not None
 
-    mnemonics = extract_mnemonics_from_commands(result.stdout)
+    captured = capfd.readouterr()
+    mnemonics = extract_mnemonics_from_commands(captured.out)
 
     wallet_status, message = verify_wallet_dir(
         wallet_path,
@@ -372,8 +373,8 @@ def test_wallet_regen(wallet_setup):
             "--mnemonic",
             mnemonics["coldkey"],
             "--no-use-password",
+            "--overwrite"
         ],
-        inputs=["y", "y"],
     )
 
     # Wait a bit to ensure file system updates modification time
@@ -412,8 +413,8 @@ def test_wallet_regen(wallet_setup):
             wallet_path,
             "--ss58-address",
             ss58_address,
+            "--overwrite"
         ],
-        inputs=["y"],
     )
 
     # Wait a bit to ensure file system updates modification time
@@ -447,8 +448,8 @@ def test_wallet_regen(wallet_setup):
             "--mnemonic",
             mnemonics["hotkey"],
             "--no-use-password",
+            "--overwrite",
         ],
-        inputs=["y"],
     )
 
     # Wait a bit to ensure file system updates modification time