module names consistent with other safe_*

gakonst · Jul 21, 2022 · 75c2eb4 · 75c2eb4
1 parent 7326d29
commit 75c2eb4
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 14 deletions.
diff --git a/ethers-contract/ethers-contract-abigen/src/contract.rs b/ethers-contract/ethers-contract-abigen/src/contract.rs
@@ -101,8 +101,7 @@ impl Context {
     /// Expands the whole rust contract
     pub fn expand(&self) -> Result<ExpandedContract> {
         let name = &self.contract_ident;
-        let name_mod =
-            util::ident(&format!("{}_mod", self.contract_ident.to_string().to_lowercase()));
+        let name_mod = util::ident(&util::safe_module_name(&self.contract_name));
         let abi_name = self.inline_abi_ident();
 
         // 0. Imports

diff --git a/ethers-contract/ethers-contract-abigen/src/multi.rs b/ethers-contract/ethers-contract-abigen/src/multi.rs
@@ -1079,7 +1079,7 @@ contract Enum {
         assert!(multi_file_mod.exists());
         let content = fs::read_to_string(&multi_file_mod).unwrap();
         assert!(content.contains("pub mod shared_types"));
-        assert!(content.contains("pub mod _enum"));
+        assert!(content.contains("pub mod enum_"));
 
         let greeter1 = multi_file_dir.join("greeter_1.rs");
         assert!(greeter1.exists());
@@ -1093,9 +1093,9 @@ contract Enum {
         assert!(!content.contains("pub struct Inner"));
         assert!(!content.contains("pub struct Stuff"));
 
-        let _enum = multi_file_dir.join("_enum.rs");
-        assert!(_enum.exists());
-        let content = fs::read_to_string(&_enum).unwrap();
+        let enum_ = multi_file_dir.join("enum_.rs");
+        assert!(enum_.exists());
+        let content = fs::read_to_string(&enum_).unwrap();
         assert!(!content.contains("pub enum Operation"));
 
         let shared_types = multi_file_dir.join("shared_types.rs");
@@ -1104,4 +1104,65 @@ contract Enum {
         assert!(content.contains("pub struct Inner"));
         assert!(content.contains("pub struct Stuff"));
     }
+
+    #[test]
+    fn can_sanitize_bindings() {
+        let tmp = TempProject::dapptools().unwrap();
+
+        tmp.add_source(
+            "Greeter",
+            r#"
+// SPDX-License-Identifier: MIT
+pragma solidity >=0.8.0;
+
+contract Mod {
+    function greet() public pure returns (uint256) {
+        return 1;
+    }
+}
+
+// from a gnosis contract
+contract Enum {
+    enum Operation {Call, DelegateCall}
+}
+"#,
+        )
+        .unwrap();
+
+        let _ = tmp.compile().unwrap();
+
+        let gen = MultiAbigen::from_json_files(tmp.artifacts_path()).unwrap();
+        let bindings = gen.build().unwrap();
+        let single_file_dir = tmp.root().join("single_bindings");
+        bindings.write_to_module(&single_file_dir, true).unwrap();
+
+        let single_file_mod = single_file_dir.join("mod.rs");
+        assert!(single_file_mod.exists());
+        let content = fs::read_to_string(&single_file_mod).unwrap();
+        println!("{}", content);
+        assert!(content.contains("pub mod mod_ {"));
+        assert!(content.contains("pub mod enum_ {"));
+
+        // multiple files
+        let gen = MultiAbigen::from_json_files(tmp.artifacts_path()).unwrap();
+        let bindings = gen.build().unwrap();
+        let multi_file_dir = tmp.root().join("multi_bindings");
+        bindings.write_to_module(&multi_file_dir, false).unwrap();
+        let multi_file_mod = multi_file_dir.join("mod.rs");
+        assert!(multi_file_mod.exists());
+        let content = fs::read_to_string(&multi_file_mod).unwrap();
+        assert!(content.contains("pub mod enum_"));
+        assert!(content.contains("pub mod mod_"));
+
+        let enum_ = multi_file_dir.join("enum_.rs");
+        assert!(enum_.exists());
+        let content = fs::read_to_string(&enum_).unwrap();
+        assert!(!content.contains("pub enum Operation"));
+
+        let mod_ = multi_file_dir.join("mod_.rs");
+        assert!(mod_.exists());
+        let content = fs::read_to_string(&mod_).unwrap();
+        println!("{}", content);
+        assert!(!content.contains("pub enum Operation"));
+    }
 }
diff --git a/ethers-contract/ethers-contract-abigen/src/util.rs b/ethers-contract/ethers-contract-abigen/src/util.rs
@@ -42,11 +42,8 @@ fn safe_identifier_name(name: String) -> String {
 
 /// converts invalid rust module names to valid ones
 pub fn safe_module_name(name: &str) -> String {
-    match safe_snake_case(name).as_ref() {
-        // handle reserve words used in contracts (eg Enum is a gnosis contract)
-        name @ ("enum" | "mod" | "module") => format!("_{}", name),
-        name => name.to_string(),
-    }
+    // handle reserve words used in contracts (eg Enum is a gnosis contract)
+    safe_ident(&safe_snake_case(name)).to_string()
 }
 
 /// Expands an identifier as snakecase and preserve any leading or trailing underscores
@@ -241,9 +238,8 @@ mod tests {
     #[test]
     fn test_safe_module_name() {
         assert_eq!(safe_module_name("Valid"), "valid");
-        assert_eq!(safe_module_name("Enum"), "_enum");
-        assert_eq!(safe_module_name("Mod"), "_mod");
-        assert_eq!(safe_module_name("Module"), "_module");
+        assert_eq!(safe_module_name("Enum"), "enum_");
+        assert_eq!(safe_module_name("Mod"), "mod_");
         assert_eq!(safe_module_name("2Two"), "_2_two");
     }
 }