diff --git a/crates/cheatcodes/assets/cheatcodes.json b/crates/cheatcodes/assets/cheatcodes.json index 7a02c6ad64fa..c8b5d3d2affe 100644 --- a/crates/cheatcodes/assets/cheatcodes.json +++ b/crates/cheatcodes/assets/cheatcodes.json @@ -6231,6 +6231,26 @@ "status": "stable", "safety": "safe" }, + { + "func": { + "id": "promptAddress", + "description": "Prompts the user for an address in the terminal.", + "declaration": "function promptAddress(string calldata promptText) external returns (address);", + "visibility": "external", + "mutability": "", + "signature": "promptAddress(string)", + "selector": "0x62ee05f4", + "selectorBytes": [ + 98, + 238, + 5, + 244 + ] + }, + "group": "filesystem", + "status": "stable", + "safety": "safe" + }, { "func": { "id": "promptSecret", @@ -6251,6 +6271,26 @@ "status": "stable", "safety": "safe" }, + { + "func": { + "id": "promptUint", + "description": "Prompts the user for uint256 in the terminal.", + "declaration": "function promptUint(string calldata promptText) external returns (uint256);", + "visibility": "external", + "mutability": "", + "signature": "promptUint(string)", + "selector": "0x652fd489", + "selectorBytes": [ + 101, + 47, + 212, + 137 + ] + }, + "group": "filesystem", + "status": "stable", + "safety": "safe" + }, { "func": { "id": "readCallers", diff --git a/crates/cheatcodes/spec/src/vm.rs b/crates/cheatcodes/spec/src/vm.rs index ae93475b2bf3..03a2cd090b85 100644 --- a/crates/cheatcodes/spec/src/vm.rs +++ b/crates/cheatcodes/spec/src/vm.rs @@ -1482,6 +1482,14 @@ interface Vm { #[cheatcode(group = Filesystem)] function promptSecret(string calldata promptText) external returns (string memory input); + /// Prompts the user for an address in the terminal. + #[cheatcode(group = Filesystem)] + function promptAddress(string calldata promptText) external returns (address); + + /// Prompts the user for uint256 in the terminal. + #[cheatcode(group = Filesystem)] + function promptUint(string calldata promptText) external returns (uint256); + // ======== Environment Variables ======== /// Sets environment variables. diff --git a/crates/cheatcodes/src/fs.rs b/crates/cheatcodes/src/fs.rs index 3ba07b5b6cd2..9b22c4d8d2ae 100644 --- a/crates/cheatcodes/src/fs.rs +++ b/crates/cheatcodes/src/fs.rs @@ -1,6 +1,8 @@ //! Implementations of [`Filesystem`](crate::Group::Filesystem) cheatcodes. +use super::string::parse; use crate::{Cheatcode, Cheatcodes, Result, Vm::*}; +use alloy_dyn_abi::DynSolType; use alloy_json_abi::ContractObject; use alloy_primitives::U256; use alloy_sol_types::SolValue; @@ -426,6 +428,20 @@ impl Cheatcode for promptSecretCall { } } +impl Cheatcode for promptAddressCall { + fn apply(&self, state: &mut Cheatcodes) -> Result { + let Self { promptText: text } = self; + parse(&prompt(state, text, prompt_input)?, &DynSolType::Address) + } +} + +impl Cheatcode for promptUintCall { + fn apply(&self, state: &mut Cheatcodes) -> Result { + let Self { promptText: text } = self; + parse(&prompt(state, text, prompt_input)?, &DynSolType::Uint(256)) + } +} + pub(super) fn write_file(state: &Cheatcodes, path: &Path, contents: &[u8]) -> Result { let path = state.config.ensure_path_allowed(path, FsAccessKind::Write)?; // write access to foundry.toml is not allowed diff --git a/testdata/cheats/Vm.sol b/testdata/cheats/Vm.sol index 1d0b6118b1a3..44f0d52ac423 100644 --- a/testdata/cheats/Vm.sol +++ b/testdata/cheats/Vm.sol @@ -307,7 +307,9 @@ interface Vm { function prevrandao(bytes32 newPrevrandao) external; function projectRoot() external view returns (string memory path); function prompt(string calldata promptText) external returns (string memory input); + function promptAddress(string calldata promptText) external returns (address); function promptSecret(string calldata promptText) external returns (string memory input); + function promptUint(string calldata promptText) external returns (uint256); function readCallers() external returns (CallerMode callerMode, address msgSender, address txOrigin); function readDir(string calldata path) external view returns (DirEntry[] memory entries); function readDir(string calldata path, uint64 maxDepth) external view returns (DirEntry[] memory entries); diff --git a/testdata/default/cheats/Prompt.t.sol b/testdata/default/cheats/Prompt.t.sol index dadfd30a97b2..9e461c2b527e 100644 --- a/testdata/default/cheats/Prompt.t.sol +++ b/testdata/default/cheats/Prompt.t.sol @@ -3,6 +3,7 @@ pragma solidity 0.8.18; import "ds-test/test.sol"; import "cheats/Vm.sol"; +import "../logs/console.sol"; contract PromptTest is DSTest { Vm constant vm = Vm(HEVM_ADDRESS); @@ -15,4 +16,14 @@ contract PromptTest is DSTest { vm._expectCheatcodeRevert(); vm.promptSecret("test"); } + + function testPrompt_Address() public { + vm._expectCheatcodeRevert(); + address test = vm.promptAddress("test"); + } + + function testPrompt_Uint() public { + vm._expectCheatcodeRevert(); + uint256 test = vm.promptUint("test"); + } }