[v10] RootClaim update

bittensor/core/async_subtensor.py

@@ -26,6 +26,7 @@
     ProxyInfo,
     ProxyConstants,
     ProxyType,
+    RootClaimType,
     SelectiveMetagraphIndex,
     SimSwapResult,
     StakeInfo,
@@ -3475,7 +3476,7 @@ async def get_root_claim_type(
         block: Optional[int] = None,
         block_hash: Optional[str] = None,
         reuse_block: bool = False,
-    ) -> str:
+    ) -> Union[str, dict]:
         """Retrieves the root claim type for a given coldkey address.
 
         Parameters:
@@ -3485,7 +3486,8 @@ async def get_root_claim_type(
             reuse_block: Whether to reuse the last-used block hash. Do not set if using block_hash or block.
 
         Returns:
-            RootClaimType value in string representation. Could be `Swap` or `Keep`.
+            Union[str, dict]: RootClaimType value. Returns string for "Swap" or "Keep",
+            or dict for "KeepSubnets" in format {"KeepSubnets": {"subnets": [1, 2, 3]}}.
         """
         block_hash = await self.determine_block_hash(block, block_hash, reuse_block)
         query = await self.substrate.query(
@@ -3495,7 +3497,55 @@ async def get_root_claim_type(
             block_hash=block_hash,
             reuse_block_hash=reuse_block,
         )
-        return next(iter(query.keys()))
+        # Query returns enum as dict: {"Swap": ()} or {"Keep": ()} or {"KeepSubnets": {"subnets": [1, 2, 3]}}
+        variant_name = next(iter(query.keys()))
+        variant_value = query[variant_name]
+
+        # For simple variants (Swap, Keep), value is empty tuple, return string
+        if not variant_value or variant_value == ():
+            return variant_name
+
+        # For KeepSubnets, value contains the data, return full dict structure
+        if isinstance(variant_value, dict) and "subnets" in variant_value:
+            subnets_raw = variant_value["subnets"]
+            subnets = list(subnets_raw[0])
+
+            return {variant_name: {"subnets": subnets}}
+
+        return {variant_name: variant_value}
+
+    async def get_root_alpha_dividends_per_subnet(
+        self,
+        hotkey_ss58: str,
+        netuid: int,
+        block: Optional[int] = None,
+        block_hash: Optional[str] = None,
+        reuse_block: bool = False,
+    ) -> Balance:
+        """Retrieves the root alpha dividends per subnet for a given hotkey.
+
+        This storage tracks the root alpha dividends that a hotkey has received on a specific subnet.
+        It is updated during block emission distribution when root alpha is distributed to validators.
+
+        Parameters:
+            hotkey_ss58: The ss58 address of the root validator hotkey.
+            netuid: The unique identifier of the subnet.
+            block: The block number to query. Do not specify if using block_hash or reuse_block.
+            block_hash: The block hash at which to check the parameter. Do not set if using block or reuse_block.
+            reuse_block: Whether to reuse the last-used block hash. Do not set if using block_hash or block.
+
+        Returns:
+            Balance: The root alpha dividends for this hotkey on this subnet in Rao, with unit set to netuid.
+        """
+        block_hash = await self.determine_block_hash(block, block_hash, reuse_block)
+        query = await self.substrate.query(
+            module="SubtensorModule",
+            storage_function="RootAlphaDividendsPerSubnet",
+            params=[netuid, hotkey_ss58],
+            block_hash=block_hash,
+            reuse_block_hash=reuse_block,
+        )
+        return Balance.from_rao(query.value).set_unit(netuid=netuid)
 
     async def get_root_claimable_rate(
         self,
@@ -7205,7 +7255,7 @@ async def set_delegate_take(
     async def set_root_claim_type(
         self,
         wallet: "Wallet",
-        new_root_claim_type: Literal["Swap", "Keep"],
+        new_root_claim_type: "Literal['Swap', 'Keep'] | RootClaimType | dict",
         period: Optional[int] = DEFAULT_PERIOD,
         raise_error: bool = False,
         wait_for_inclusion: bool = True,
@@ -7215,7 +7265,11 @@ async def set_root_claim_type(
 
         Parameters:
             wallet: Bittensor Wallet instance.
-            new_root_claim_type: The new root claim type to set. Could be either "Swap" or "Keep".
+            new_root_claim_type: The new root claim type to set. Can be:
+                - String: "Swap" or "Keep"
+                - RootClaimType: RootClaimType.Swap, RootClaimType.Keep
+                - Dict: {"KeepSubnets": {"subnets": [1, 2, 3]}}
+                - Callable: RootClaimType.KeepSubnets([1, 2, 3])
             period: The number of blocks during which the transaction will remain valid after it's submitted. If the
                 transaction is not included in a block within that number of blocks, it will expire and be rejected. You
                 can think of it as an expiration date for the transaction.

bittensor/core/chain_data/__init__.py

@@ -24,6 +24,7 @@
 from .prometheus_info import PrometheusInfo
 from .proposal_vote_data import ProposalVoteData
 from .proxy import ProxyConstants, ProxyInfo, ProxyType, ProxyAnnouncementInfo
+from .root_claim import RootClaimType
 from .scheduled_coldkey_swap_info import ScheduledColdkeySwapInfo
 from .stake_info import StakeInfo
 from .sim_swap import SimSwapResult
@@ -59,6 +60,7 @@
     "ProxyAnnouncementInfo",
     "ProxyInfo",
     "ProxyType",
+    "RootClaimType",
     "ScheduledColdkeySwapInfo",
     "SelectiveMetagraphIndex",
     "SimSwapResult",

bittensor/core/chain_data/proxy.py

@@ -79,6 +79,7 @@ class ProxyType(str, Enum):
 
         RootClaim: Allows only root claim operations. Permitted operations:
             - claim_root
+            - set_root_claim_type
 
     Note:
         The values match exactly with the ProxyType enum defined in the Subtensor runtime. Any changes to the runtime

bittensor/core/chain_data/root_claim.py

@@ -0,0 +1,122 @@
+from enum import Enum
+from typing import Literal, Callable
+from dataclasses import dataclass
+
+
+@dataclass
+class KeepSubnetsDescriptor:
+    """Descriptor that allows callable syntax for KeepSubnets variant."""
+
+    subnets: list[int]
+
+    def __post_init__(self):
+        """Validate subnets after initialization."""
+        if not self.subnets:
+            raise ValueError(
+                "KeepSubnets must have at least one subnet represented by a netuid in the list."
+            )
+        if not all(isinstance(netuid, int) for netuid in self.subnets):
+            raise ValueError("All subnet IDs must be integers in the list.")
+
+    def to_dict(self) -> dict:
+        """Converts the descriptor to the required dictionary format."""
+        return {"KeepSubnets": {"subnets": self.subnets}}
+
+    def __call__(self, subnets: list[int]) -> dict:
+        """Allows calling the descriptor with subnets to create a new instance and return dict."""
+        return KeepSubnetsDescriptor(subnets).to_dict()
+
+    def __get__(self, instance, owner) -> Callable[[list[int]], dict]:
+        """Descriptor protocol - returns a callable that creates the dict."""
+
+        def create(subnets: list[int]) -> dict:
+            return KeepSubnetsDescriptor(subnets).to_dict()
+
+        return create
+
+
+class RootClaimType(str, Enum):
+    """
+    Enumeration of root claim types in the Bittensor network.
+
+    This enum defines how coldkeys manage their root alpha emissions:
+    - Swap: Swap any alpha emission for TAO
+    - Keep: Keep all alpha emission
+    - KeepSubnets: Keep alpha emission for specified subnets, swap everything else
+
+    The values match exactly with the RootClaimTypeEnum defined in the Subtensor runtime.
+    """
+
+    Swap = "Swap"
+    Keep = "Keep"
+    KeepSubnets = KeepSubnetsDescriptor
+
+    @classmethod
+    def normalize(
+        cls, value: "Literal['Swap', 'Keep'] | RootClaimType | dict"
+    ) -> str | dict:
+        """
+        Normalizes a root claim type to a format suitable for Substrate calls.
+
+        This method handles various input formats:
+        - String values ("Swap", "Keep") → returns string
+        - Enum values (RootClaimType.Swap) → returns string
+        - Dict values ({"KeepSubnets": {"subnets": [1, 2, 3]}}) → returns dict as-is
+        - Callable KeepSubnets([1, 2, 3]) → returns dict
+
+        Parameters:
+            value: The root claim type in any supported format.
+
+        Returns:
+            Normalized value - string for Swap/Keep or dict for KeepSubnets.
+
+        Raises:
+            ValueError: If the value is not a valid root claim type or KeepSubnets has no subnets.
+            TypeError: If the value type is not supported.
+        """
+        # Handle KeepSubnetsDescriptor instance
+        if isinstance(value, KeepSubnetsDescriptor):
+            return value.to_dict()
+
+        # Handle enum instance
+        if isinstance(value, RootClaimType):
+            # If it's KeepSubnets, it's actually the descriptor, so this shouldn't happen
+            # But if someone accesses it directly, we need to handle it
+            if value == "KeepSubnets":
+                raise ValueError(
+                    "KeepSubnets must be called with subnet list: RootClaimType.KeepSubnets([1, 2, 3])"
+                )
+            return value.value
+
+        # Handle string values
+        if isinstance(value, str):
+            if value in ("Swap", "Keep"):
+                return value
+            elif value == "KeepSubnets":
+                raise ValueError(
+                    "KeepSubnets must be provided as dict or called: RootClaimType.KeepSubnets([1, 2, 3])"
+                )
+            else:
+                raise ValueError(
+                    f"Invalid root claim type: {value}. "
+                    f"Valid types are: 'Swap', 'Keep', or KeepSubnets dict/callable"
+                )
+
+        # Handle dict values (for KeepSubnets)
+        if isinstance(value, dict):
+            if "KeepSubnets" in value and isinstance(value["KeepSubnets"], dict):
+                subnets = value["KeepSubnets"].get("subnets", [])
+                if not subnets:
+                    raise ValueError("KeepSubnets must have at least one subnet")
+                if not all(isinstance(netuid, int) for netuid in subnets):
+                    raise ValueError("All subnet IDs must be integers")
+                return value
+            else:
+                raise ValueError(
+                    f"Invalid dict format for root claim type. "
+                    f"Expected {{'KeepSubnets': {{'subnets': [1, 2, 3]}}}}, got {value}"
+                )
+
+        raise TypeError(
+            f"root_claim_type must be str, RootClaimType, or dict, got {type(value).__name__}"
+        )

bittensor/core/extrinsics/asyncex/root.py

@@ -1,6 +1,7 @@
 import asyncio
 from typing import Optional, TYPE_CHECKING, Literal
 
+from bittensor.core.chain_data import RootClaimType
 from bittensor.core.extrinsics.pallets import SubtensorModule
 from bittensor.core.types import ExtrinsicResponse
 from bittensor.utils import u16_normalized_float
@@ -147,7 +148,7 @@ async def root_register_extrinsic(
 async def set_root_claim_type_extrinsic(
     subtensor: "AsyncSubtensor",
     wallet: "Wallet",
-    new_root_claim_type: Literal["Swap", "Keep"],
+    new_root_claim_type: "Literal['Swap', 'Keep'] | RootClaimType | dict",
     period: Optional[int] = None,
     raise_error: bool = False,
     wait_for_inclusion: bool = True,
@@ -158,7 +159,11 @@ async def set_root_claim_type_extrinsic(
     Parameters:
         subtensor: Subtensor instance to interact with the blockchain.
         wallet: Bittensor Wallet instance.
-        new_root_claim_type: The new root claim type to set. Could be either "Swap" or "Keep".
+        new_root_claim_type: The new root claim type to set. Can be:
+            - String: "Swap" or "Keep"
+            - RootClaimType: RootClaimType.Swap, RootClaimType.Keep
+            - Dict: {"KeepSubnets": {"subnets": [1, 2, 3]}}
+            - Callable: RootClaimType.KeepSubnets([1, 2, 3])
         period: The number of blocks during which the transaction will remain valid after it's submitted. If the
             transaction is not included in a block within that number of blocks, it will expire and be rejected. You can
             think of it as an expiration date for the transaction.
@@ -175,8 +180,10 @@ async def set_root_claim_type_extrinsic(
         ).success:
             return unlocked
 
+        normalized_type = RootClaimType.normalize(new_root_claim_type)
+
         call = await SubtensorModule(subtensor).set_root_claim_type(
-            new_root_claim_type=new_root_claim_type
+            new_root_claim_type=normalized_type
         )
 
         return await subtensor.sign_and_send_extrinsic(

bittensor/core/extrinsics/pallets/subtensor_module.py

@@ -511,12 +511,14 @@ def set_pending_childkey_cooldown(
 
     def set_root_claim_type(
         self,
-        new_root_claim_type: Literal["Swap", "Keep"],
+        new_root_claim_type: Literal["Swap", "Keep"] | dict,
     ) -> Call:
         """Returns GenericCall instance for Subtensor function SubtensorModule.set_root_claim_type.
 
         Parameters:
-            new_root_claim_type: The new root claim type.
+            new_root_claim_type: The new root claim type. Can be:
+                - String: "Swap" or "Keep"
+                - Dict: {"KeepSubnets": {"subnets": [1, 2, 3]}}
 
         Returns:
             GenericCall instance.

bittensor/core/extrinsics/root.py

@@ -1,6 +1,7 @@
 import time
 from typing import Literal, Optional, TYPE_CHECKING
 
+from bittensor.core.chain_data import RootClaimType
 from bittensor.core.extrinsics.pallets import SubtensorModule
 from bittensor.core.types import ExtrinsicResponse, UIDs
 from bittensor.utils import u16_normalized_float
@@ -142,7 +143,7 @@ def root_register_extrinsic(
 def set_root_claim_type_extrinsic(
     subtensor: "Subtensor",
     wallet: "Wallet",
-    new_root_claim_type: Literal["Swap", "Keep"],
+    new_root_claim_type: "Literal['Swap', 'Keep'] | RootClaimType | dict",
     period: Optional[int] = None,
     raise_error: bool = False,
     wait_for_inclusion: bool = True,
@@ -153,7 +154,11 @@ def set_root_claim_type_extrinsic(
     Parameters:
         subtensor: Subtensor instance to interact with the blockchain.
         wallet: Bittensor Wallet instance.
-        new_root_claim_type: The new root claim type to set. Could be either "Swap" or "Keep".
+        new_root_claim_type: The new root claim type to set. Can be:
+            - String: "Swap" or "Keep"
+            - RootClaimType: RootClaimType.Swap, RootClaimType.Keep
+            - Dict: {"KeepSubnets": {"subnets": [1, 2, 3]}}
+            - Callable: RootClaimType.KeepSubnets([1, 2, 3])
         period: The number of blocks during which the transaction will remain valid after it's submitted. If the
             transaction is not included in a block within that number of blocks, it will expire and be rejected. You can
             think of it as an expiration date for the transaction.
@@ -170,8 +175,10 @@ def set_root_claim_type_extrinsic(
         ).success:
             return unlocked
 
+        normalized_type = RootClaimType.normalize(new_root_claim_type)
+
         call = SubtensorModule(subtensor).set_root_claim_type(
-            new_root_claim_type=new_root_claim_type
+            new_root_claim_type=normalized_type
         )
 
         return subtensor.sign_and_send_extrinsic(