refactor: post review

src/cli.rs

@@ -8,14 +8,14 @@ pub struct Args {
     pub solc: String,
 
     /// Number of handler to generate
-    #[arg(short = 'n', long, default_value_t = 2)]
+    #[arg(short = 'n', long, default_value_t = 2, value_parser = clap::value_parser!(u8).range(1..))]
     pub nb_handlers: u8,
 
     /// Number of properties contract to generate
-    #[arg(short = 'p', long, default_value_t = 2)]
+    #[arg(short = 'p', long, default_value_t = 2, value_parser = clap::value_parser!(u8).range(1..))]
     pub nb_properties: u8,
 
     /// Overwrite existing files
-    #[arg(short, long, default_value = "false")]
+    #[arg(short, long, action = clap::ArgAction::SetFalse)]
     pub overwrite: bool,
 }

src/gen.rs

@@ -28,45 +28,53 @@ fn parse_parents(parents: &[Contract]) -> String {
         .to_string()
 }
 
-/// Generate the parents contracts and write them to the temp folder
-fn generate_parents(
-    contract_type: ContractType,
-    args: &Args,
-    path: &Path,
-) -> Result<Vec<Contract>> {
-    let nb_parents = match contract_type {
-        ContractType::Handler => args.nb_handlers,
-        ContractType::Property => args.nb_properties,
-        ContractType::EntryPoint | ContractType::Setup => {
-            return Err(anyhow::anyhow!("Invalid contract type in generate parents"))?
-        }
-    };
-
-    let mut parents = Vec::new();
+/// create a vec of contracts of a given type
+fn create_contracts(contract_type: ContractType, count: u8, path: &Path) -> Result<Vec<Contract>> {
+    let mut contracts = Vec::new();
 
+    // directories
     DirBuilder::new()
         .recursive(true)
         .create(path)
         .context(format!(
-            "Fail to create folder for {}",
+            "Failed to create directory for {} contracts",
             contract_type.directory_name()
         ))?;
 
-    for i in 0..nb_parents {
+    for i in 0..count {
         let contract = ContractBuilder::new()
             .with_type(&contract_type)
             .with_name(format!("{}{}", contract_type.name(), (b'A' + i) as char))
             .build();
 
         contract.write_rendered_contract(path).context(format!(
-            "Failed to write rendered {} parent",
+            "Failed to write rendered {} contract",
             contract_type.name()
         ))?;
 
-        parents.push(contract);
+        contracts.push(contract);
     }
 
-    Ok(parents)
+    Ok(contracts)
+}
+
+/// Generate parents contracts and write them to a temp folder
+fn generate_parents(
+    contract_type: ContractType,
+    args: &Args,
+    path: &Path,
+) -> Result<Vec<Contract>> {
+    // Determine the number of parents to generate
+    let count = match contract_type {
+        ContractType::Handler => args.nb_handlers,
+        ContractType::Property => args.nb_properties,
+        _ => {
+            return Err(anyhow::anyhow!("Invalid contract type in generate_parents"));
+        }
+    };
+
+    // Use the helper function to generate the contracts
+    create_contracts(contract_type, count, path)
 }
 
 /// Move the content of a temp folder to the fuzz test folder

src/types.rs

@@ -51,16 +51,6 @@ impl ContractBuilder {
         }
     }
 
-    pub fn with_licence(mut self, licence: String) -> Self {
-        self.licence = licence;
-        self
-    }
-
-    pub fn with_solc(mut self, solc: String) -> Self {
-        self.solc = solc;
-        self
-    }
-
     pub fn with_imports(mut self, imports: String) -> Self {
         self.imports = imports;
         self
@@ -77,9 +67,9 @@ impl ContractBuilder {
     }
 
     pub fn with_type(mut self, contract_type: &ContractType) -> Self {
-        self.imports = contract_type.import();
-        self.name = contract_type.name();
-        self.parents = contract_type.import_name();
+        self.imports = contract_type.import().to_owned();
+        self.name = contract_type.name().to_owned();
+        self.parents = contract_type.import_name().to_owned();
         self
     }
 
@@ -104,42 +94,42 @@ pub enum ContractType {
 
 /// Hold the contract type specific information
 impl ContractType {
-    pub fn directory_name(&self) -> String {
+    pub fn directory_name(&self) -> &'static str {
         match self {
-            ContractType::Handler => "handlers".to_string(),
-            ContractType::Property => "properties".to_string(),
-            _ => "".to_string(),
+            ContractType::Handler => "handlers",
+            ContractType::Property => "properties",
+            _ => "",
         }
     }
 
-    pub fn name(&self) -> String {
+    pub fn name(&self) -> &'static str {
         match self {
-            ContractType::Handler => "Handlers".to_string(),
-            ContractType::Property => "Properties".to_string(),
-            ContractType::EntryPoint => "FuzzTest".to_string(),
-            ContractType::Setup => "Setup".to_string(),
+            ContractType::Handler => "Handlers",
+            ContractType::Property => "Properties",
+            ContractType::EntryPoint => "FuzzTest",
+            ContractType::Setup => "Setup",
         }
     }
 
-    pub fn import(&self) -> String {
+    pub fn import(&self) -> &'static str {
         match self {
-            ContractType::Handler => "import {Setup} from '../Setup.t.sol';".to_string(),
+            ContractType::Handler => "import {Setup} from '../Setup.t.sol';",
             ContractType::Property => {
-                "import {HandlersParent} from '../handlers/HandlersParent.t.sol';".to_string()
+                "import {HandlersParent} from '../handlers/HandlersParent.t.sol';"
             }
             ContractType::EntryPoint => {
-                "import {PropertiesParent} from './properties/PropertiesParent.t.sol';".to_string()
+                "import {PropertiesParent} from './properties/PropertiesParent.t.sol';"
             }
-            ContractType::Setup => "".to_string(),
+            ContractType::Setup => "",
         }
     }
 
-    pub fn import_name(&self) -> String {
+    pub fn import_name(&self) -> &'static str {
         match self {
-            ContractType::Handler => "Setup".to_string(),
-            ContractType::Property => "HandlersParent".to_string(),
-            ContractType::EntryPoint => "PropertiesParent".to_string(),
-            ContractType::Setup => "".to_string(),
+            ContractType::Handler => "Setup",
+            ContractType::Property => "HandlersParent",
+            ContractType::EntryPoint => "PropertiesParent",
+            ContractType::Setup => "",
         }
     }
 }