Feat/Root claim update for netuids

bittensor_cli/cli.py

@@ -7216,6 +7216,12 @@ def stake_set_claim_type(
             None,
             help="Claim type: 'keep' or 'swap'. If not provided, you'll be prompted to choose.",
         ),
+        netuids: Optional[str] = typer.Option(
+            None,
+            "--netuids",
+            "-n",
+            help="Netuids to select. Supports ranges and comma-separated values, e.g., '1-5,10,20-30'.",
+        ),
         wallet_name: Optional[str] = Options.wallet_name,
         wallet_path: Optional[str] = Options.wallet_path,
         wallet_hotkey: Optional[str] = Options.wallet_hotkey,
@@ -7233,29 +7239,22 @@ def stake_set_claim_type(
         [bold]Claim Types:[/bold]
         • [green]Swap[/green]: Future Root Alpha Emissions are swapped to TAO and added to root stake (default)
         • [yellow]Keep[/yellow]: Future Root Alpha Emissions are kept as Alpha tokens
+        • [cyan]Keep Specific[/cyan]: Keep specific subnets as Alpha, swap others to TAO
 
         USAGE:
 
         [green]$[/green] btcli stake claim
         [green]$[/green] btcli stake claim keep
         [green]$[/green] btcli stake claim swap
+        [green]$[/green] btcli stake claim keep --netuids 1-5,10,20-30
+        [green]$[/green] btcli stake claim swap --netuids 1-30
 
         With specific wallet:
 
         [green]$[/green] btcli stake claim swap --wallet-name my_wallet
         """
         self.verbosity_handler(quiet, verbose, json_output)
 
-        if claim_type is not None:
-            claim_type_normalized = claim_type.capitalize()
-            if claim_type_normalized not in ["Keep", "Swap"]:
-                err_console.print(
-                    f":cross_mark: [red]Invalid claim type '{claim_type}'. Must be 'keep' or 'swap'.[/red]"
-                )
-                raise typer.Exit()
-        else:
-            claim_type_normalized = None
-
         wallet = self.wallet_ask(
             wallet_name,
             wallet_path,
@@ -7266,7 +7265,8 @@ def stake_set_claim_type(
             claim_stake.set_claim_type(
                 wallet=wallet,
                 subtensor=self.initialize_chain(network),
-                claim_type=claim_type_normalized,
+                claim_type=claim_type,
+                netuids=netuids,
                 prompt=prompt,
                 json_output=json_output,
             )

bittensor_cli/src/bittensor/subtensor_interface.py

@@ -652,6 +652,26 @@ async def subnet_exists(
         )
         return result
 
+    async def total_networks(
+        self, block_hash: Optional[str] = None, reuse_block: bool = False
+    ) -> int:
+        """
+        Returns the total number of subnets in the Bittensor network.
+
+        :param block_hash: The hash of the blockchain block number at which to check the subnet existence.
+        :param reuse_block: Whether to reuse the last-used block hash.
+
+        :return: The total number of subnets in the network.
+        """
+        result = await self.query(
+            module="SubtensorModule",
+            storage_function="TotalNetworks",
+            params=[],
+            block_hash=block_hash,
+            reuse_block_hash=reuse_block,
+        )
+        return result
+
     async def get_subnet_state(
         self, netuid: int, block_hash: Optional[str] = None
     ) -> Optional["SubnetState"]:
@@ -1832,21 +1852,25 @@ async def get_coldkey_claim_type(
         coldkey_ss58: str,
         block_hash: Optional[str] = None,
         reuse_block: bool = False,
-    ) -> str:
+    ) -> dict:
         """
         Retrieves the root claim type for a specific coldkey.
 
         Root claim types control how staking emissions are handled on the ROOT network (subnet 0):
         - "Swap": Future Root Alpha Emissions are swapped to TAO at claim time and added to your root stake
         - "Keep": Future Root Alpha Emissions are kept as Alpha
+        - "KeepSubnets": Specific subnets kept as Alpha, rest swapped to TAO
 
         Args:
             coldkey_ss58: The SS58 address of the coldkey to query.
             block_hash: The hash of the blockchain block number for the query.
             reuse_block: Whether to reuse the last-used blockchain block hash.
 
         Returns:
-            str: The root claim type for the coldkey ("Swap" or "Keep").
+            dict: Claim type information in one of these formats:
+                - {"type": "Swap"}
+                - {"type": "Keep"}
+                - {"type": "KeepSubnets", "subnets": [1, 5, 10, ...]}
         """
         result = await self.query(
             module="SubtensorModule",
@@ -1857,14 +1881,22 @@ async def get_coldkey_claim_type(
         )
 
         if result is None:
-            return "Swap"
-        return next(iter(result.keys()))
+            return {"type": "Swap"}
+
+        claim_type_key = next(iter(result.keys()))
+
+        if claim_type_key == "KeepSubnets":
+            subnets_data = result["KeepSubnets"]["subnets"]
+            subnet_list = sorted([subnet for subnet in subnets_data[0]])
+            return {"type": "KeepSubnets", "subnets": subnet_list}
+        else:
+            return {"type": claim_type_key}
 
     async def get_all_coldkeys_claim_type(
         self,
         block_hash: Optional[str] = None,
         reuse_block: bool = False,
-    ) -> dict[str, str]:
+    ) -> dict[str, dict]:
         """
         Retrieves all root claim types for all coldkeys in the network.
 
@@ -1873,7 +1905,7 @@ async def get_all_coldkeys_claim_type(
             reuse_block: Whether to reuse the last-used blockchain block hash.
 
         Returns:
-            dict[str, str]: A dictionary mapping coldkey SS58 addresses to their root claim type ("Keep" or "Swap").
+            dict[str, dict]: Mapping of coldkey SS58 addresses to claim type dicts
         """
         result = await self.substrate.query_map(
             module="SubtensorModule",
@@ -1884,10 +1916,20 @@ async def get_all_coldkeys_claim_type(
         )
 
         root_claim_types = {}
-        async for coldkey, claim_type in result:
+        async for coldkey, claim_type_data in result:
             coldkey_ss58 = decode_account_id(coldkey[0])
-            claim_type = next(iter(claim_type.value.keys()))
-            root_claim_types[coldkey_ss58] = claim_type
+
+            claim_type_key = next(iter(claim_type_data.value.keys()))
+
+            if claim_type_key == "KeepSubnets":
+                subnets_data = claim_type_data.value["KeepSubnets"]["subnets"]
+                subnet_list = sorted([subnet for subnet in subnets_data[0]])
+                root_claim_types[coldkey_ss58] = {
+                    "type": "KeepSubnets",
+                    "subnets": subnet_list,
+                }
+            else:
+                root_claim_types[coldkey_ss58] = {"type": claim_type_key}
 
         return root_claim_types
 

bittensor_cli/src/bittensor/utils.py

@@ -1051,6 +1051,58 @@ def group_subnets(registrations):
     return ", ".join(ranges)
 
 
+def parse_subnet_range(input_str: str, total_subnets: int) -> list[int]:
+    """
+    Parse subnet range input like "1-24, 30-40, 5".
+
+    Args:
+        input_str: Comma-separated list of subnets and ranges
+                  Examples: "1-5", "1,2,3", "1-5, 10, 20-25"
+        total_subnets: Total number of subnets available
+
+    Returns:
+        Sorted list of unique subnet IDs
+
+    Raises:
+        ValueError: If input format is invalid
+
+    Examples:
+        >>> parse_subnet_range("1-5, 10")
+        [1, 2, 3, 4, 5, 10]
+        >>> parse_subnet_range("5, 3, 1")
+        [1, 3, 5]
+    """
+    subnets = set()
+    parts = [p.strip() for p in input_str.split(",") if p.strip()]
+    for part in parts:
+        if "-" in part:
+            try:
+                start, end = part.split("-", 1)
+                start_num = int(start.strip())
+                end_num = int(end.strip())
+
+                if start_num > end_num:
+                    raise ValueError(f"Invalid range '{part}': start must be ≤ end")
+
+                if end_num - start_num > total_subnets:
+                    raise ValueError(
+                        f"Range '{part}' is not valid (total of {total_subnets} subnets)"
+                    )
+
+                subnets.update(range(start_num, end_num + 1))
+            except ValueError as e:
+                if "invalid literal" in str(e):
+                    raise ValueError(f"Invalid range '{part}': must be 'start-end'")
+                raise
+        else:
+            try:
+                subnets.add(int(part))
+            except ValueError:
+                raise ValueError(f"Invalid subnet ID '{part}': must be a number")
+
+    return sorted(subnets)
+
+
 def validate_chain_endpoint(endpoint_url) -> tuple[bool, str]:
     parsed = urlparse(endpoint_url)
     if parsed.scheme not in ("ws", "wss"):