Port get_next_address

Author: Quexington

chia/_tests/cmds/cmd_test_utils.py

@@ -34,6 +34,8 @@
 from chia.wallet.util.tx_config import CoinSelectionConfig, TXConfig
 from chia.wallet.util.wallet_types import WalletType
 from chia.wallet.wallet_request_types import (
+    GetNextAddress,
+    GetNextAddressResponse,
     GetSyncStatusResponse,
     GetTransaction,
     GetTransactionResponse,
@@ -250,13 +252,13 @@ async def get_spendable_coins(
         unconfirmed_additions = [Coin(bytes32([7] * 32), bytes32([8] * 32), uint64(1234580000))]
         return confirmed_records, unconfirmed_removals, unconfirmed_additions
 
-    async def get_next_address(self, wallet_id: int, new_address: bool) -> str:
-        self.add_to_log("get_next_address", (wallet_id, new_address))
+    async def get_next_address(self, request: GetNextAddress) -> GetNextAddressResponse:
+        self.add_to_log("get_next_address", (request.wallet_id, request.new_address))
         addr = encode_puzzle_hash(bytes32([self.wallet_index] * 32), "xch")
         self.wallet_index += 1
         if self.wallet_index > 254:
             self.wallet_index = 1
-        return addr
+        return GetNextAddressResponse(request.wallet_id, addr)
 
     async def send_transaction_multi(
         self,

chia/_tests/cmds/wallet/test_wallet.py

@@ -47,6 +47,8 @@
     CreateOfferForIDsResponse,
     FungibleAsset,
     GetHeightInfoResponse,
+    GetNextAddress,
+    GetNextAddressResponse,
     GetTransaction,
     GetTransactions,
     GetTransactionsResponse,
@@ -493,11 +495,11 @@ def test_get_address(capsys: object, get_test_cli_clients: tuple[TestRpcClients,
 
     # set RPC Client
     class GetAddressWalletRpcClient(TestWalletRpcClient):
-        async def get_next_address(self, wallet_id: int, new_address: bool) -> str:
-            self.add_to_log("get_next_address", (wallet_id, new_address))
-            if new_address:
-                return encode_puzzle_hash(get_bytes32(3), "xch")
-            return encode_puzzle_hash(get_bytes32(4), "xch")
+        async def get_next_address(self, request: GetNextAddress) -> GetNextAddressResponse:
+            self.add_to_log("get_next_address", (request.wallet_id, request.new_address))
+            if request.new_address:
+                return GetNextAddressResponse(request.wallet_id, encode_puzzle_hash(get_bytes32(3), "xch"))
+            return GetNextAddressResponse(request.wallet_id, encode_puzzle_hash(get_bytes32(4), "xch"))
 
     inst_rpc_client = GetAddressWalletRpcClient()
     test_rpc_clients.wallet_rpc_client = inst_rpc_client

chia/_tests/wallet/rpc/test_wallet_rpc.py

@@ -118,6 +118,7 @@
     DIDTransferDID,
     DIDUpdateMetadata,
     FungibleAsset,
+    GetNextAddress,
     GetNotifications,
     GetPrivateKey,
     GetSyncStatusResponse,
@@ -195,11 +196,11 @@ async def farm_transaction(
 
 
 async def generate_funds(full_node_api: FullNodeSimulator, wallet_bundle: WalletBundle, num_blocks: int = 1) -> int:
-    wallet_id = 1
-    initial_balances = (
-        await wallet_bundle.rpc_client.get_wallet_balance(GetWalletBalance(uint32(wallet_id)))
-    ).wallet_balance
-    ph: bytes32 = decode_puzzle_hash(await wallet_bundle.rpc_client.get_next_address(wallet_id, True))
+    wallet_id = uint32(1)
+    initial_balances = (await wallet_bundle.rpc_client.get_wallet_balance(GetWalletBalance(wallet_id))).wallet_balance
+    ph: bytes32 = decode_puzzle_hash(
+        (await wallet_bundle.rpc_client.get_next_address(GetNextAddress(wallet_id, True))).address
+    )
     generated_funds = 0
     for _ in range(num_blocks):
         await full_node_api.farm_new_transaction_block(FarmNewBlockProtocol(ph))
@@ -1213,8 +1214,8 @@ async def test_cat_endpoints(wallet_environments: WalletTestFramework, wallet_ty
         ]
     )
 
-    addr_0 = await env_0.rpc_client.get_next_address(cat_0_id, False)
-    addr_1 = await env_1.rpc_client.get_next_address(cat_1_id, False)
+    addr_0 = (await env_0.rpc_client.get_next_address(GetNextAddress(cat_0_id, False))).address
+    addr_1 = (await env_1.rpc_client.get_next_address(GetNextAddress(cat_1_id, False))).address
 
     assert addr_0 != addr_1
 
@@ -1420,7 +1421,7 @@ async def test_offer_endpoints(wallet_environments: WalletTestFramework, wallet_
 
     # Creates a wallet for the same CAT on wallet_2 and send 4 CAT from wallet_1 to it
     await env_2.rpc_client.create_wallet_for_existing_cat(cat_asset_id)
-    wallet_2_address = await env_2.rpc_client.get_next_address(cat_wallet_id, False)
+    wallet_2_address = (await env_2.rpc_client.get_next_address(GetNextAddress(cat_wallet_id, False))).address
     adds = [{"puzzle_hash": decode_puzzle_hash(wallet_2_address), "amount": uint64(4), "memos": ["the cat memo"]}]
     tx_res = (
         await env_1.rpc_client.send_transaction_multi(
@@ -2149,7 +2150,7 @@ async def test_key_and_address_endpoints(wallet_rpc_environment: WalletRpcTestEn
     wallet_node: WalletNode = env.wallet_1.node
     client: WalletRpcClient = env.wallet_1.rpc_client
 
-    address = await client.get_next_address(1, True)
+    address = (await client.get_next_address(GetNextAddress(uint32(1), True))).address
     assert len(address) > 10
 
     pks = (await client.get_public_keys()).pk_fingerprints
@@ -2793,7 +2794,7 @@ async def test_set_wallet_resync_on_startup(wallet_rpc_environment: WalletRpcTes
 
     nft_wallet = await wc.create_new_nft_wallet(None)
     nft_wallet_id = nft_wallet["wallet_id"]
-    address = await wc.get_next_address(env.wallet_1.wallet.id(), True)
+    address = (await wc.get_next_address(GetNextAddress(env.wallet_1.wallet.id(), True))).address
     await wc.mint_nft(
         request=NFTMintNFTRequest(
             wallet_id=nft_wallet_id,

chia/cmds/wallet_funcs.py

@@ -53,6 +53,7 @@
     DIDTransferDID,
     DIDUpdateMetadata,
     FungibleAsset,
+    GetNextAddress,
     GetNotifications,
     GetTransaction,
     GetTransactions,
@@ -416,7 +417,7 @@ async def get_address(
     root_path: pathlib.Path, wallet_rpc_port: Optional[int], fp: Optional[int], wallet_id: int, new_address: bool
 ) -> None:
     async with get_wallet_client(root_path, wallet_rpc_port, fp) as (wallet_client, _, _):
-        res = await wallet_client.get_next_address(wallet_id, new_address)
+        res = (await wallet_client.get_next_address(GetNextAddress(uint32(wallet_id), new_address))).address
         print(res)
 
 

chia/wallet/wallet_request_types.py

@@ -412,6 +412,21 @@ class GetTransactionCountResponse(Streamable):
     count: uint16
 
 
+@streamable
+@dataclass(frozen=True)
+class GetNextAddress(Streamable):
+    wallet_id: uint32
+    new_address: bool = False
+    save_derivations: bool = True
+
+
+@streamable
+@dataclass(frozen=True)
+class GetNextAddressResponse(Streamable):
+    wallet_id: uint32
+    address: str
+
+
 @streamable
 @dataclass(frozen=True)
 class GetOffersCountResponse(Streamable):

chia/wallet/wallet_rpc_api.py

@@ -170,6 +170,8 @@
     GenerateMnemonicResponse,
     GetHeightInfoResponse,
     GetLoggedInFingerprintResponse,
+    GetNextAddress,
+    GetNextAddressResponse,
     GetNotifications,
     GetNotificationsResponse,
     GetPrivateKey,
@@ -1550,33 +1552,29 @@ async def get_transaction_count(self, request: GetTransactionCount) -> GetTransa
             uint16(count),
         )
 
-    async def get_next_address(self, request: dict[str, Any]) -> EndpointResult:
+    @marshal
+    async def get_next_address(self, request: GetNextAddress) -> GetNextAddressResponse:
         """
         Returns a new address
         """
-        if request["new_address"] is True:
-            create_new = True
-        else:
-            create_new = False
-        wallet_id = uint32(request["wallet_id"])
-        wallet = self.service.wallet_state_manager.wallets[wallet_id]
+        wallet = self.service.wallet_state_manager.wallets[request.wallet_id]
         selected = self.service.config["selected_network"]
         prefix = self.service.config["network_overrides"]["config"][selected]["address_prefix"]
         if wallet.type() in {WalletType.STANDARD_WALLET, WalletType.CAT, WalletType.CRCAT, WalletType.RCAT}:
             async with self.service.wallet_state_manager.new_action_scope(
-                DEFAULT_TX_CONFIG, push=request.get("save_derivations", True)
+                DEFAULT_TX_CONFIG, push=request.save_derivations
             ) as action_scope:
                 raw_puzzle_hash = await action_scope.get_puzzle_hash(
-                    self.service.wallet_state_manager, override_reuse_puzhash_with=not create_new
+                    self.service.wallet_state_manager, override_reuse_puzhash_with=not request.new_address
                 )
             address = encode_puzzle_hash(raw_puzzle_hash, prefix)
         else:
             raise ValueError(f"Wallet type {wallet.type()} cannot create puzzle hashes")
 
-        return {
-            "wallet_id": wallet_id,
-            "address": address,
-        }
+        return GetNextAddressResponse(
+            request.wallet_id,
+            address,
+        )
 
     @tx_endpoint(push=True)
     async def send_transaction(

chia/wallet/wallet_rpc_client.py

@@ -87,6 +87,8 @@
     GetCATListResponse,
     GetHeightInfoResponse,
     GetLoggedInFingerprintResponse,
+    GetNextAddress,
+    GetNextAddressResponse,
     GetNotifications,
     GetNotificationsResponse,
     GetOffersCountResponse,
@@ -281,11 +283,8 @@ async def get_transaction_count(self, request: GetTransactionCount) -> GetTransa
             await self.fetch("get_transaction_count", request.to_json_dict())
         )
 
-    async def get_next_address(self, wallet_id: int, new_address: bool) -> str:
-        request = {"wallet_id": wallet_id, "new_address": new_address}
-        response = await self.fetch("get_next_address", request)
-        # TODO: casting due to lack of type checked deserialization
-        return cast(str, response["address"])
+    async def get_next_address(self, request: GetNextAddress) -> GetNextAddressResponse:
+        return GetNextAddressResponse.from_json_dict(await self.fetch("get_next_address", request.to_json_dict()))
 
     async def send_transaction(
         self,