meshtensor-cli 9.18.1__py3-none-any.whl

meshtensor_cli/src/commands/stake/auto_staking.py

@@ -0,0 +1,306 @@
+import asyncio
+import json
+from typing import Optional, TYPE_CHECKING
+
+from meshtensor_wallet import Wallet
+from rich import box
+from rich.table import Table
+
+from meshtensor_cli.src import COLOR_PALETTE
+from meshtensor_cli.src.meshtensor.utils import (
+    confirm_action,
+    console,
+    json_console,
+    print_success,
+    get_subnet_name,
+    is_valid_ss58_address,
+    print_error,
+    unlock_key,
+    print_extrinsic_id,
+)
+
+if TYPE_CHECKING:
+    from meshtensor_cli.src.meshtensor.meshtensor_interface import MeshtensorInterface
+
+
+async def show_auto_stake_destinations(
+    wallet: Optional[Wallet],
+    meshtensor: "MeshtensorInterface",
+    coldkey_ss58: Optional[str] = None,
+    json_output: bool = False,
+) -> Optional[dict[int, dict[str, Optional[str]]]]:
+    """Display auto-stake destinations for the supplied wallet."""
+
+    wallet_name: Optional[str] = wallet.name if wallet else None
+    coldkey_ss58 = coldkey_ss58 or (wallet.coldkeypub.ss58_address if wallet else None)
+    if not coldkey_ss58:
+        raise ValueError("A wallet or coldkey SS58 address must be provided")
+
+    with console.status(
+        f"Retrieving auto-stake configuration from {meshtensor.network}...",
+        spinner="earth",
+    ):
+        chain_head = await meshtensor.substrate.get_chain_head()
+        (
+            subnet_info,
+            auto_destinations,
+            identities,
+            delegate_identities,
+        ) = await asyncio.gather(
+            meshtensor.all_subnets(block_hash=chain_head),
+            meshtensor.get_auto_stake_destinations(
+                coldkey_ss58=coldkey_ss58,
+                block_hash=chain_head,
+                reuse_block=True,
+            ),
+            meshtensor.fetch_coldkey_hotkey_identities(block_hash=chain_head),
+            meshtensor.get_delegate_identities(block_hash=chain_head),
+        )
+
+    subnet_map = {info.netuid: info for info in subnet_info}
+    auto_destinations = auto_destinations or {}
+    identities = identities or {}
+    delegate_identities = delegate_identities or {}
+    hotkey_identities = identities.get("hotkeys", {})
+
+    def resolve_identity(hotkey: str) -> Optional[str]:
+        if not hotkey:
+            return None
+
+        identity_entry = hotkey_identities.get(hotkey, {}).get("identity")
+        if identity_entry:
+            display_name = identity_entry.get("name") or identity_entry.get("display")
+            if display_name:
+                return display_name
+
+        delegate_info = delegate_identities.get(hotkey)
+        if delegate_info and getattr(delegate_info, "display", ""):
+            return delegate_info.display
+
+        return None
+
+    coldkey_display = wallet_name
+    if not coldkey_display:
+        coldkey_identity = identities.get("coldkeys", {}).get(coldkey_ss58, {})
+        if identity_data := coldkey_identity.get("identity"):
+            coldkey_display = identity_data.get("name") or identity_data.get("display")
+        if not coldkey_display:
+            coldkey_display = f"{coldkey_ss58[:6]}...{coldkey_ss58[-6:]}"
+
+    rows = []
+    data_output: dict[int, dict[str, Optional[str]]] = {}
+
+    for netuid in sorted(subnet_map):
+        subnet = subnet_map[netuid]
+        subnet_name = get_subnet_name(subnet)
+        hotkey_ss58 = auto_destinations.get(netuid)
+        identity_str = resolve_identity(hotkey_ss58) if hotkey_ss58 else None
+        is_custom = hotkey_ss58 is not None
+
+        data_output[netuid] = {
+            "subnet_name": subnet_name,
+            "status": "custom" if is_custom else "default",
+            "destination": hotkey_ss58,
+            "identity": identity_str,
+        }
+
+        if json_output:
+            continue
+
+        status_text = (
+            f"[{COLOR_PALETTE['STAKE']['STAKE_ALPHA']}]Custom[/{COLOR_PALETTE['STAKE']['STAKE_ALPHA']}]"
+            if is_custom
+            else f"[{COLOR_PALETTE['GENERAL']['HINT']}]Default[/{COLOR_PALETTE['GENERAL']['HINT']}]"
+        )
+
+        rows.append(
+            (
+                str(netuid),
+                subnet_name,
+                status_text,
+                hotkey_ss58,
+                identity_str or "",
+            )
+        )
+
+    if json_output:
+        json_console.print(json.dumps(data_output))
+        return data_output
+
+    table = Table(
+        title=(
+            f"\n[{COLOR_PALETTE['GENERAL']['HEADER']}]Auto Stake Destinations"
+            f" for [bold]{coldkey_display}[/bold]\n"
+            f"Network: {meshtensor.network}\n"
+            f"Coldkey: {coldkey_ss58}\n"
+            f"[/{COLOR_PALETTE['GENERAL']['HEADER']}]"
+        ),
+        show_edge=False,
+        header_style="bold white",
+        border_style="bright_black",
+        style="bold",
+        title_justify="center",
+        show_lines=False,
+        pad_edge=True,
+        box=box.SIMPLE_HEAD,
+    )
+
+    table.add_column(
+        "Netuid", style=COLOR_PALETTE["GENERAL"]["SYMBOL"], justify="center"
+    )
+    table.add_column("Subnet", style="cyan", justify="left")
+    table.add_column("Status", style="white", justify="center")
+    table.add_column(
+        "Destination Hotkey", style=COLOR_PALETTE["GENERAL"]["HOTKEY"], justify="center"
+    )
+    table.add_column(
+        "Identity", style=COLOR_PALETTE["GENERAL"]["SUBHEADING"], justify="left"
+    )
+
+    for row in rows:
+        table.add_row(*row)
+
+    console.print(table)
+    console.print(
+        f"\n[{COLOR_PALETTE['GENERAL']['SUBHEADING']}]Total subnets:[/] {len(subnet_map)}  "
+        f"[{COLOR_PALETTE['GENERAL']['SUBHEADING']}]Custom destinations:[/] {len(auto_destinations)}"
+    )
+
+    return None
+
+
+async def set_auto_stake_destination(
+    wallet: Wallet,
+    meshtensor: "MeshtensorInterface",
+    netuid: int,
+    hotkey_ss58: str,
+    proxy: Optional[str] = None,
+    wait_for_inclusion: bool = True,
+    wait_for_finalization: bool = False,
+    prompt_user: bool = True,
+    decline: bool = False,
+    quiet: bool = False,
+    json_output: bool = False,
+) -> bool:
+    """Set the auto-stake destination hotkey for a coldkey on a subnet."""
+
+    if not is_valid_ss58_address(hotkey_ss58):
+        print_error("You entered an invalid hotkey ss58 address")
+        return False
+
+    try:
+        chain_head = await meshtensor.substrate.get_chain_head()
+        subnet_info, identities, delegate_identities = await asyncio.gather(
+            meshtensor.subnet(netuid, block_hash=chain_head),
+            meshtensor.fetch_coldkey_hotkey_identities(block_hash=chain_head),
+            meshtensor.get_delegate_identities(block_hash=chain_head),
+        )
+    except ValueError:
+        print_error(f"Subnet with netuid {netuid} does not exist")
+        return False
+
+    hotkey_identity = ""
+    identities = identities or {}
+    delegate_identities = delegate_identities or {}
+
+    hotkey_identity_entry = identities.get("hotkeys", {}).get(hotkey_ss58, {})
+    if identity_data := hotkey_identity_entry.get("identity"):
+        hotkey_identity = (
+            identity_data.get("name") or identity_data.get("display") or ""
+        )
+    if not hotkey_identity:
+        delegate_info = delegate_identities.get(hotkey_ss58)
+        if delegate_info and getattr(delegate_info, "display", ""):
+            hotkey_identity = delegate_info.display
+
+    if prompt_user and not json_output:
+        table = Table(
+            title=(
+                f"\n[{COLOR_PALETTE['GENERAL']['HEADER']}]Confirm Auto-Stake Destination"
+                f"[/{COLOR_PALETTE['GENERAL']['HEADER']}]"
+            ),
+            show_edge=False,
+            header_style="bold white",
+            border_style="bright_black",
+            style="bold",
+            title_justify="center",
+            show_lines=False,
+            pad_edge=True,
+            box=box.SIMPLE_HEAD,
+        )
+        table.add_column(
+            "Netuid", justify="center", style=COLOR_PALETTE["GENERAL"]["SYMBOL"]
+        )
+        table.add_column("Subnet", style="cyan", justify="left")
+        table.add_column(
+            "Destination Hotkey",
+            style=COLOR_PALETTE["GENERAL"]["HOTKEY"],
+            justify="center",
+        )
+        table.add_column(
+            "Identity", style=COLOR_PALETTE["GENERAL"]["SUBHEADING"], justify="left"
+        )
+        table.add_row(
+            str(netuid),
+            get_subnet_name(subnet_info),
+            hotkey_ss58,
+            hotkey_identity or "",
+        )
+        console.print(table)
+
+        if not confirm_action(
+            "\nSet this auto-stake destination?",
+            default=True,
+            decline=decline,
+            quiet=quiet,
+        ):
+            return False
+
+    if not unlock_key(wallet).success:
+        return False
+
+    call = await meshtensor.substrate.compose_call(
+        call_module="MeshtensorModule",
+        call_function="set_coldkey_auto_stake_hotkey",
+        call_params={
+            "netuid": netuid,
+            "hotkey": hotkey_ss58,
+        },
+    )
+
+    with console.status(
+        f":satellite: Setting auto-stake destination on [white]{meshtensor.network}[/white]...",
+        spinner="earth",
+    ):
+        success, error_message, ext_receipt = await meshtensor.sign_and_send_extrinsic(
+            call,
+            wallet,
+            wait_for_inclusion=wait_for_inclusion,
+            wait_for_finalization=wait_for_finalization,
+            proxy=proxy,
+        )
+
+    ext_id = await ext_receipt.get_extrinsic_identifier() if success else None
+
+    if json_output:
+        json_console.print(
+            json.dumps(
+                {
+                    "success": success,
+                    "error": error_message,
+                    "netuid": netuid,
+                    "hotkey": hotkey_ss58,
+                    "extrinsic_identifier": ext_id,
+                }
+            )
+        )
+
+    if success:
+        await print_extrinsic_id(ext_receipt)
+        print_success(
+            f"[dark_sea_green3]Auto-stake destination set for netuid {netuid}[/dark_sea_green3]"
+        )
+        return True
+
+    print_error(f"Failed: {error_message}")
+    return False