Port extend_derivation_index

Author: Quexington

chia/_tests/wallet/test_wallet_state_manager.py

@@ -21,7 +21,7 @@
 from chia.wallet.transaction_record import TransactionRecord
 from chia.wallet.util.transaction_type import TransactionType
 from chia.wallet.util.wallet_types import WalletType
-from chia.wallet.wallet_request_types import PushTransactions
+from chia.wallet.wallet_request_types import ExtendDerivationIndex, PushTransactions
 from chia.wallet.wallet_rpc_api import MAX_DERIVATION_INDEX_DELTA
 from chia.wallet.wallet_spend_bundle import WalletSpendBundle
 from chia.wallet.wallet_state_manager import WalletStateManager
@@ -420,7 +420,7 @@ async def get_puzzle_hash_state() -> PuzzleHashState:
             (0,),
         )
     with pytest.raises(ValueError):
-        await rpc_client.extend_derivation_index(0)
+        await rpc_client.extend_derivation_index(ExtendDerivationIndex(uint32(0)))
 
     # Reset to a normal state
     await wsm.puzzle_store.delete_wallet(wsm.main_wallet.id())
@@ -431,15 +431,17 @@ async def get_puzzle_hash_state() -> PuzzleHashState:
 
     # Test an index already created
     with pytest.raises(ValueError):
-        await rpc_client.extend_derivation_index(0)
+        await rpc_client.extend_derivation_index(ExtendDerivationIndex(uint32(0)))
 
     # Test an index too far in the future
     with pytest.raises(ValueError):
-        await rpc_client.extend_derivation_index(MAX_DERIVATION_INDEX_DELTA + expected_state.highest_index + 1)
+        await rpc_client.extend_derivation_index(
+            ExtendDerivationIndex(uint32(MAX_DERIVATION_INDEX_DELTA + expected_state.highest_index + 1))
+        )
 
     # Test the actual functionality
-    assert await rpc_client.extend_derivation_index(expected_state.highest_index + 5) == str(
-        expected_state.highest_index + 5
-    )
+    assert (
+        await rpc_client.extend_derivation_index(ExtendDerivationIndex(uint32(expected_state.highest_index + 5)))
+    ).index == expected_state.highest_index + 5
     expected_state = PuzzleHashState(expected_state.highest_index + 5, expected_state.used_up_to_index)
     assert await get_puzzle_hash_state() == expected_state

chia/cmds/wallet_funcs.py

@@ -53,6 +53,7 @@
     DIDSetWalletName,
     DIDTransferDID,
     DIDUpdateMetadata,
+    ExtendDerivationIndex,
     FungibleAsset,
     GetNextAddress,
     GetNotifications,
@@ -443,8 +444,8 @@ async def update_derivation_index(
 ) -> None:
     async with get_wallet_client(root_path, wallet_rpc_port, fp) as (wallet_client, _, _):
         print("Updating derivation index... This may take a while.")
-        res = await wallet_client.extend_derivation_index(index)
-        print(f"Updated derivation index: {res}")
+        res = await wallet_client.extend_derivation_index(ExtendDerivationIndex(uint32(index)))
+        print(f"Updated derivation index: {res.index}")
         print("Your balances may take a while to update.")
 
 

chia/wallet/wallet_request_types.py

@@ -444,6 +444,18 @@ class GetCurrentDerivationIndexResponse(Streamable):
     index: Optional[uint32]
 
 
+@streamable
+@dataclass(frozen=True)
+class ExtendDerivationIndex(Streamable):
+    index: uint32
+
+
+@streamable
+@dataclass(frozen=True)
+class ExtendDerivationIndexResponse(Streamable):
+    index: Optional[uint32]
+
+
 @streamable
 @dataclass(frozen=True)
 class GetOffersCountResponse(Streamable):

chia/wallet/wallet_rpc_api.py

@@ -166,6 +166,8 @@
     Empty,
     ExecuteSigningInstructions,
     ExecuteSigningInstructionsResponse,
+    ExtendDerivationIndex,
+    ExtendDerivationIndexResponse,
     GatherSigningInfo,
     GatherSigningInfoResponse,
     GenerateMnemonicResponse,
@@ -1879,30 +1881,26 @@ async def get_current_derivation_index(self, request: Empty) -> GetCurrentDeriva
 
         return GetCurrentDerivationIndexResponse(index)
 
-    async def extend_derivation_index(self, request: dict[str, Any]) -> dict[str, Any]:
+    @marshal
+    async def extend_derivation_index(self, request: ExtendDerivationIndex) -> ExtendDerivationIndexResponse:
         assert self.service.wallet_state_manager is not None
 
-        # Require a new max derivation index
-        if "index" not in request:
-            raise ValueError("Derivation index is required")
-
         # Require that the wallet is fully synced
         synced = await self.service.wallet_state_manager.synced()
         if synced is False:
             raise ValueError("Wallet needs to be fully synced before extending derivation index")
 
-        index = uint32(request["index"])
         current: Optional[uint32] = await self.service.wallet_state_manager.puzzle_store.get_last_derivation_path()
 
         # Additional sanity check that the wallet is synced
         if current is None:
             raise ValueError("No current derivation record found, unable to extend index")
 
         # Require that the new index is greater than the current index
-        if index <= current:
+        if request.index <= current:
             raise ValueError(f"New derivation index must be greater than current index: {current}")
 
-        if index - current > MAX_DERIVATION_INDEX_DELTA:
+        if request.index - current > MAX_DERIVATION_INDEX_DELTA:
             raise ValueError(
                 "Too many derivations requested. "
                 f"Use a derivation index less than {current + MAX_DERIVATION_INDEX_DELTA + 1}"
@@ -1912,14 +1910,13 @@ async def extend_derivation_index(self, request: dict[str, Any]) -> dict[str, An
         # to preserve the current last used index, so we call create_more_puzzle_hashes with
         # mark_existing_as_used=False
         result = await self.service.wallet_state_manager.create_more_puzzle_hashes(
-            from_zero=False, mark_existing_as_used=False, up_to_index=index, num_additional_phs=0
+            from_zero=False, mark_existing_as_used=False, up_to_index=request.index, num_additional_phs=0
         )
         await result.commit(self.service.wallet_state_manager)
 
-        updated: Optional[uint32] = await self.service.wallet_state_manager.puzzle_store.get_last_derivation_path()
-        updated_index = updated if updated is not None else None
+        updated_index = await self.service.wallet_state_manager.puzzle_store.get_last_derivation_path()
 
-        return {"success": True, "index": updated_index}
+        return ExtendDerivationIndexResponse(updated_index)
 
     @marshal
     async def get_notifications(self, request: GetNotifications) -> GetNotificationsResponse:

chia/wallet/wallet_rpc_client.py

@@ -82,6 +82,8 @@
     DLUpdateRootResponse,
     ExecuteSigningInstructions,
     ExecuteSigningInstructionsResponse,
+    ExtendDerivationIndex,
+    ExtendDerivationIndexResponse,
     GatherSigningInfo,
     GatherSigningInfoResponse,
     GenerateMnemonicResponse,
@@ -373,10 +375,10 @@ async def delete_unconfirmed_transactions(self, request: DeleteUnconfirmedTransa
     async def get_current_derivation_index(self) -> GetCurrentDerivationIndexResponse:
         return GetCurrentDerivationIndexResponse.from_json_dict(await self.fetch("get_current_derivation_index", {}))
 
-    async def extend_derivation_index(self, index: int) -> str:
-        response = await self.fetch("extend_derivation_index", {"index": index})
-        updated_index = response["index"]
-        return str(updated_index)
+    async def extend_derivation_index(self, request: ExtendDerivationIndex) -> ExtendDerivationIndexResponse:
+        return ExtendDerivationIndexResponse.from_json_dict(
+            await self.fetch("extend_derivation_index", request.to_json_dict())
+        )
 
     async def get_farmed_amount(self, include_pool_rewards: bool = False) -> dict[str, Any]:
         return await self.fetch("get_farmed_amount", {"include_pool_rewards": include_pool_rewards})