forge/cmd/generate/
mod.rs

1use clap::{Parser, Subcommand};
2use eyre::Result;
3use foundry_common::fs;
4use std::path::Path;
5use yansi::Paint;
6
7/// CLI arguments for `forge generate`.
8#[derive(Debug, Parser)]
9pub struct GenerateArgs {
10    #[command(subcommand)]
11    pub sub: GenerateSubcommands,
12}
13
14#[derive(Debug, Subcommand)]
15pub enum GenerateSubcommands {
16    /// Scaffolds test file for given contract.
17    Test(GenerateTestArgs),
18}
19
20#[derive(Debug, Parser)]
21pub struct GenerateTestArgs {
22    /// Contract name for test generation.
23    #[arg(long, short, value_name = "CONTRACT_NAME")]
24    pub contract_name: String,
25}
26
27impl GenerateTestArgs {
28    pub fn run(self) -> Result<()> {
29        let contract_name = format_identifier(&self.contract_name, true);
30        let instance_name = format_identifier(&self.contract_name, false);
31
32        // Create the test file content.
33        let test_content = include_str!("../../../assets/generated/TestTemplate.t.sol");
34        let test_content = test_content
35            .replace("{contract_name}", &contract_name)
36            .replace("{instance_name}", &instance_name);
37
38        // Create the test directory if it doesn't exist.
39        fs::create_dir_all("test")?;
40
41        // Define the test file path
42        let test_file_path = Path::new("test").join(format!("{contract_name}.t.sol"));
43
44        // Write the test content to the test file.
45        fs::write(&test_file_path, test_content)?;
46
47        sh_println!("{} test file: {}", "Generated".green(), test_file_path.to_str().unwrap())?;
48        Ok(())
49    }
50}
51
52/// Utility function to convert an identifier to pascal or camel case.
53fn format_identifier(input: &str, is_pascal_case: bool) -> String {
54    let mut result = String::new();
55    let mut capitalize_next = is_pascal_case;
56
57    for word in input.split_whitespace() {
58        if !word.is_empty() {
59            let (first, rest) = word.split_at(1);
60            let formatted_word = if capitalize_next {
61                format!("{}{}", first.to_uppercase(), rest)
62            } else {
63                format!("{}{}", first.to_lowercase(), rest)
64            };
65            capitalize_next = true;
66            result.push_str(&formatted_word);
67        }
68    }
69    result
70}