...

chia/_tests/cmds/test_click_types.py

@@ -1,12 +1,11 @@
 from __future__ import annotations
 
-from dataclasses import dataclass, field
 from decimal import Decimal
 from pathlib import Path
-from typing import Any, cast
 
+import click
 import pytest
-from click import BadParameter, Context
+from click import BadParameter
 
 from chia.cmds.cmd_classes import ChiaCliContext
 from chia.cmds.param_types import (
@@ -39,9 +38,9 @@
 overflow_decimal = Decimal(overflow_decimal_str)
 
 
-@dataclass
-class FakeContext:
-    obj: dict[Any, Any] = field(default_factory=dict)
+@click.command()
+def a_command_for_testing() -> None:
+    pass  # pragma: no cover
 
 
 def test_click_tx_fee_type() -> None:
@@ -115,11 +114,11 @@ def test_click_amount_type() -> None:
 
 
 def test_click_address_type() -> None:
+    context = click.Context(command=a_command_for_testing)
+    chia_context = ChiaCliContext.set_default(context)
     # this makes us not have to use a config file
-    context = cast(
-        Context,
-        FakeContext(obj=ChiaCliContext(expected_prefix="xch").to_click()),
-    )
+    chia_context.expected_prefix = "xch"
+
     std_cli_address = CliAddress(burn_ph, burn_address, AddressType.XCH)
     nft_cli_address = CliAddress(burn_ph, burn_nft_addr, AddressType.DID)
     # Test CliAddress (Generally is not used)
@@ -137,7 +136,7 @@ def test_click_address_type() -> None:
     # check error handling
     with pytest.raises(BadParameter):
         AddressParamType().convert("test", None, None)
-    with pytest.raises(AttributeError):  # attribute error because the context does not have a real error handler
+    with pytest.raises(click.BadParameter):
         AddressParamType().convert(burn_address_txch, None, context)
     with pytest.raises(BadParameter):
         AddressParamType().convert(burn_bad_prefix, None, None)
@@ -153,11 +152,13 @@ def test_click_address_type() -> None:
 
 
 def test_click_address_type_config(root_path_populated_with_config: Path) -> None:
+    context = click.Context(command=a_command_for_testing)
+    chia_context = ChiaCliContext.set_default(context)
     # set a root path in context.
-    context = cast(Context, FakeContext(obj=ChiaCliContext(root_path=root_path_populated_with_config).to_click()))
+    chia_context.root_path = root_path_populated_with_config
     # run test that should pass
     assert AddressParamType().convert(burn_address, None, context) == CliAddress(burn_ph, burn_address, AddressType.XCH)
-    assert ChiaCliContext.from_click(context).expected_prefix == "xch"  # validate that the prefix was set correctly
+    assert ChiaCliContext.set_default(context).expected_prefix == "xch"  # validate that the prefix was set correctly
     # use txch address
     with pytest.raises(AttributeError):  # attribute error because the context does not have a real error handler
         AddressParamType().convert(burn_address_txch, None, context)

chia/cmds/beta.py

@@ -38,7 +38,7 @@ def beta_cmd() -> None:
 @click.option("-i", "--interval", help="System metrics will be logged based on this interval", type=int, required=False)
 @click.pass_context
 def configure(ctx: click.Context, path: Optional[str], interval: Optional[int]) -> None:
-    root_path = ChiaCliContext.from_click(ctx).root_path
+    root_path = ChiaCliContext.set_default(ctx).root_path
     with lock_and_load_config(root_path, "config.yaml") as config:
         if "beta" not in config:
             raise click.ClickException("beta test mode is not enabled, enable it first with `chia beta enable`")
@@ -80,7 +80,7 @@ def configure(ctx: click.Context, path: Optional[str], interval: Optional[int])
 @click.option("-p", "--path", help="The beta mode root path", type=str, required=False)
 @click.pass_context
 def enable_cmd(ctx: click.Context, force: bool, path: Optional[str]) -> None:
-    root_path = ChiaCliContext.from_click(ctx).root_path
+    root_path = ChiaCliContext.set_default(ctx).root_path
     with lock_and_load_config(root_path, "config.yaml") as config:
         if config.get("beta", {}).get("enabled", False):
             raise click.ClickException("beta test mode is already enabled")
@@ -108,7 +108,7 @@ def enable_cmd(ctx: click.Context, force: bool, path: Optional[str]) -> None:
 @beta_cmd.command("disable", help="Disable beta test mode")
 @click.pass_context
 def disable_cmd(ctx: click.Context) -> None:
-    root_path = ChiaCliContext.from_click(ctx).root_path
+    root_path = ChiaCliContext.set_default(ctx).root_path
     with lock_and_load_config(root_path, "config.yaml") as config:
         if not config.get("beta", {}).get("enabled", False):
             raise click.ClickException("beta test mode is not enabled")
@@ -122,7 +122,7 @@ def disable_cmd(ctx: click.Context) -> None:
 @beta_cmd.command("prepare_submission", help="Prepare the collected log data for submission")
 @click.pass_context
 def prepare_submission_cmd(ctx: click.Context) -> None:
-    with lock_and_load_config(ChiaCliContext.from_click(ctx).root_path, "config.yaml") as config:
+    with lock_and_load_config(ChiaCliContext.set_default(ctx).root_path, "config.yaml") as config:
         beta_root_path = config.get("beta", {}).get("path", None)
         if beta_root_path is None:
             raise click.ClickException("beta test mode not enabled. Run `chia beta enable` first.")
@@ -174,7 +174,7 @@ def add_files(paths: list[Path]) -> int:
 @beta_cmd.command("status", help="Show the current beta configuration")
 @click.pass_context
 def status(ctx: click.Context) -> None:
-    with lock_and_load_config(ChiaCliContext.from_click(ctx).root_path, "config.yaml") as config:
+    with lock_and_load_config(ChiaCliContext.set_default(ctx).root_path, "config.yaml") as config:
         beta_config = config.get("beta")
         if beta_config is None:
             raise click.ClickException("beta test mode is not enabled, enable it first with `chia beta enable`")

chia/cmds/chia.py

@@ -64,8 +64,8 @@ def cli(
 ) -> None:
     from pathlib import Path
 
-    ctx.ensure_object(dict)
-    ctx.obj.update(ChiaCliContext(root_path=Path(root_path)).to_click())
+    context = ChiaCliContext.set_default(ctx=ctx)
+    context.root_path = Path(root_path)
 
     # keys_root_path and passphrase_file will be None if the passphrase options have been
     # scrubbed from the CLI options
@@ -117,7 +117,7 @@ def run_daemon_cmd(ctx: click.Context, wait_for_unlock: bool) -> None:
 
     wait_for_unlock = wait_for_unlock and Keychain.is_keyring_locked()
 
-    asyncio.run(async_run_daemon(ChiaCliContext.from_click(ctx).root_path, wait_for_unlock=wait_for_unlock))
+    asyncio.run(async_run_daemon(ChiaCliContext.set_default(ctx).root_path, wait_for_unlock=wait_for_unlock))
 
 
 cli.add_command(keys_cmd)

chia/cmds/cmd_classes.py

@@ -14,7 +14,6 @@
     Optional,
     Protocol,
     Union,
-    cast,
     final,
     get_args,
     get_origin,
@@ -76,16 +75,11 @@ class ChiaCliContext:
     expected_currency_prefix: Optional[str] = None
 
     @classmethod
-    def from_click(cls, ctx: click.Context) -> ChiaCliContext:
-        if ctx.obj is None:
-            # TODO: should we set it up on the ctx here?
-            return cls()
-
-        existing = cast(Optional[ChiaCliContext], ctx.obj.get(cls.context_dict_key))
-        if existing is None:
-            return cls()
-
-        return existing
+    def set_default(cls, ctx: click.Context) -> ChiaCliContext:
+        ctx.ensure_object(dict)
+        self = ctx.obj.setdefault(cls.context_dict_key, cls())
+        assert isinstance(self, cls)
+        return self
 
     def to_click(self) -> dict[str, object]:
         return {self.context_dict_key: self}
@@ -157,7 +151,7 @@ def apply_decorators(self, cmd: SyncCmd) -> SyncCmd:
 
             def strip_click_context(func: SyncCmd) -> SyncCmd:
                 def _inner(ctx: click.Context, **kwargs: Any) -> None:
-                    context = ChiaCliContext.from_click(ctx)
+                    context = ChiaCliContext.set_default(ctx)
                     func(context=context, **kwargs)
 
                 return _inner

chia/cmds/configure.py

@@ -314,7 +314,7 @@ def configure_cmd(
     seeder_nameserver: str,
 ) -> None:
     configure(
-        ChiaCliContext.from_click(ctx).root_path,
+        ChiaCliContext.set_default(ctx).root_path,
         set_farmer_peer,
         set_node_introducer,
         set_fullnode_port,

chia/cmds/dao.py

@@ -65,7 +65,7 @@ def dao_add_cmd(
 
     asyncio.run(
         add_dao_wallet(
-            ChiaCliContext.from_click(ctx).root_path, wallet_rpc_port, fingerprint, name, treasury_id, filter_amount
+            ChiaCliContext.set_default(ctx).root_path, wallet_rpc_port, fingerprint, name, treasury_id, filter_amount
         )
     )
 
@@ -189,7 +189,7 @@ def dao_create_cmd(
 
     return asyncio.run(
         create_dao_wallet(
-            ChiaCliContext.from_click(ctx).root_path,
+            ChiaCliContext.set_default(ctx).root_path,
             wallet_rpc_port,
             fingerprint,
             fee,
@@ -240,7 +240,7 @@ def dao_get_id_cmd(
 ) -> None:
     from chia.cmds.dao_funcs import get_treasury_id
 
-    asyncio.run(get_treasury_id(ChiaCliContext.from_click(ctx).root_path, wallet_rpc_port, fingerprint, wallet_id))
+    asyncio.run(get_treasury_id(ChiaCliContext.set_default(ctx).root_path, wallet_rpc_port, fingerprint, wallet_id))
 
 
 @dao_cmd.command("add_funds", short_help="Send funds to a DAO treasury", no_args_is_help=True)
@@ -291,7 +291,7 @@ def dao_add_funds_cmd(
 
     return asyncio.run(
         add_funds_to_treasury(
-            ChiaCliContext.from_click(ctx).root_path,
+            ChiaCliContext.set_default(ctx).root_path,
             wallet_rpc_port,
             fingerprint,
             wallet_id,
@@ -330,7 +330,9 @@ def dao_get_balance_cmd(
 ) -> None:
     from chia.cmds.dao_funcs import get_treasury_balance
 
-    asyncio.run(get_treasury_balance(ChiaCliContext.from_click(ctx).root_path, wallet_rpc_port, fingerprint, wallet_id))
+    asyncio.run(
+        get_treasury_balance(ChiaCliContext.set_default(ctx).root_path, wallet_rpc_port, fingerprint, wallet_id)
+    )
 
 
 @dao_cmd.command("rules", short_help="Get the current rules governing the DAO", no_args_is_help=True)
@@ -352,7 +354,7 @@ def dao_rules_cmd(
 ) -> None:
     from chia.cmds.dao_funcs import get_rules
 
-    asyncio.run(get_rules(ChiaCliContext.from_click(ctx).root_path, wallet_rpc_port, fingerprint, wallet_id))
+    asyncio.run(get_rules(ChiaCliContext.set_default(ctx).root_path, wallet_rpc_port, fingerprint, wallet_id))
 
 
 # ----------------------------------------------------------------------------------------
@@ -390,7 +392,7 @@ def dao_list_proposals_cmd(
 
     asyncio.run(
         list_proposals(
-            ChiaCliContext.from_click(ctx).root_path, wallet_rpc_port, fingerprint, wallet_id, include_closed
+            ChiaCliContext.set_default(ctx).root_path, wallet_rpc_port, fingerprint, wallet_id, include_closed
         )
     )
 
@@ -423,7 +425,7 @@ def dao_show_proposal_cmd(
     from chia.cmds.dao_funcs import show_proposal
 
     asyncio.run(
-        show_proposal(ChiaCliContext.from_click(ctx).root_path, wallet_rpc_port, fingerprint, wallet_id, proposal_id)
+        show_proposal(ChiaCliContext.set_default(ctx).root_path, wallet_rpc_port, fingerprint, wallet_id, proposal_id)
     )
 
 
@@ -488,7 +490,7 @@ def dao_vote_cmd(
 
     return asyncio.run(
         vote_on_proposal(
-            ChiaCliContext.from_click(ctx).root_path,
+            ChiaCliContext.set_default(ctx).root_path,
             wallet_rpc_port,
             fingerprint,
             wallet_id,
@@ -561,7 +563,7 @@ def dao_close_proposal_cmd(
 
     return asyncio.run(
         close_proposal(
-            ChiaCliContext.from_click(ctx).root_path,
+            ChiaCliContext.set_default(ctx).root_path,
             wallet_rpc_port,
             fingerprint,
             wallet_id,
@@ -625,7 +627,7 @@ def dao_lockup_coins_cmd(
 
     return asyncio.run(
         lockup_coins(
-            ChiaCliContext.from_click(ctx).root_path,
+            ChiaCliContext.set_default(ctx).root_path,
             wallet_rpc_port,
             fingerprint,
             wallet_id,
@@ -676,7 +678,7 @@ def dao_release_coins_cmd(
 
     return asyncio.run(
         release_coins(
-            ChiaCliContext.from_click(ctx).root_path,
+            ChiaCliContext.set_default(ctx).root_path,
             wallet_rpc_port,
             fingerprint,
             wallet_id,
@@ -726,7 +728,7 @@ def dao_exit_lockup_cmd(
 
     return asyncio.run(
         exit_lockup(
-            ChiaCliContext.from_click(ctx).root_path,
+            ChiaCliContext.set_default(ctx).root_path,
             wallet_rpc_port,
             fingerprint,
             wallet_id,
@@ -830,7 +832,7 @@ def dao_create_spend_proposal_cmd(
 
     return asyncio.run(
         create_spend_proposal(
-            ChiaCliContext.from_click(ctx).root_path,
+            ChiaCliContext.set_default(ctx).root_path,
             wallet_rpc_port,
             fingerprint,
             wallet_id,
@@ -942,7 +944,7 @@ def dao_create_update_proposal_cmd(
 
     return asyncio.run(
         create_update_proposal(
-            ChiaCliContext.from_click(ctx).root_path,
+            ChiaCliContext.set_default(ctx).root_path,
             wallet_rpc_port,
             fingerprint,
             wallet_id,
@@ -1025,7 +1027,7 @@ def dao_create_mint_proposal_cmd(
 
     return asyncio.run(
         create_mint_proposal(
-            ChiaCliContext.from_click(ctx).root_path,
+            ChiaCliContext.set_default(ctx).root_path,
             wallet_rpc_port,
             fingerprint,
             wallet_id,

chia/cmds/db.py

@@ -42,7 +42,7 @@ def db_upgrade_cmd(
 ) -> None:
     try:
         db_upgrade_func(
-            ChiaCliContext.from_click(ctx).root_path,
+            ChiaCliContext.set_default(ctx).root_path,
             None if in_db_path is None else Path(in_db_path),
             None if out_db_path is None else Path(out_db_path),
             no_update_config=no_update_config,
@@ -64,7 +64,7 @@ def db_upgrade_cmd(
 def db_validate_cmd(ctx: click.Context, in_db_path: Optional[str], validate_blocks: bool) -> None:
     try:
         db_validate_func(
-            ChiaCliContext.from_click(ctx).root_path,
+            ChiaCliContext.set_default(ctx).root_path,
             None if in_db_path is None else Path(in_db_path),
             validate_blocks=validate_blocks,
         )
@@ -79,7 +79,7 @@ def db_validate_cmd(ctx: click.Context, in_db_path: Optional[str], validate_bloc
 def db_backup_cmd(ctx: click.Context, db_backup_file: Optional[str], no_indexes: bool) -> None:
     try:
         db_backup_func(
-            ChiaCliContext.from_click(ctx).root_path,
+            ChiaCliContext.set_default(ctx).root_path,
             None if db_backup_file is None else Path(db_backup_file),
             no_indexes=no_indexes,
         )

chia/cmds/farm.py

@@ -69,7 +69,7 @@ def summary_cmd(
             wallet_rpc_port,
             harvester_rpc_port,
             farmer_rpc_port,
-            root_path=ChiaCliContext.from_click(ctx).root_path,
+            root_path=ChiaCliContext.set_default(ctx).root_path,
         )
     )
 
@@ -97,4 +97,4 @@ def challenges_cmd(ctx: click.Context, farmer_rpc_port: Optional[int], limit: in
 
     from chia.cmds.farm_funcs import challenges
 
-    asyncio.run(challenges(ChiaCliContext.from_click(ctx).root_path, farmer_rpc_port, limit))
+    asyncio.run(challenges(ChiaCliContext.set_default(ctx).root_path, farmer_rpc_port, limit))

chia/cmds/init.py

@@ -56,7 +56,7 @@ def init_cmd(
 
     init(
         Path(create_certs) if create_certs is not None else None,
-        ChiaCliContext.from_click(ctx).root_path,
+        ChiaCliContext.set_default(ctx).root_path,
         fix_ssl_permissions,
         testnet,
         v1_db,