use TokenStream in place of str

opentensor · Nov 13, 2024 · 63e9715 · 63e9715
1 parent 1d6c938
commit 63e9715
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 32 deletions.
diff --git a/support/linting/src/forbid_as_primitive.rs b/support/linting/src/forbid_as_primitive.rs
@@ -45,10 +45,11 @@ fn is_as_primitive(ident: &Ident) -> bool {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use quote::quote;
 
-    fn lint(input: &str) -> Result {
-        let expr: ExprMethodCall = syn::parse_str(input).expect("should only use on a method call");
+    fn lint(input: proc_macro2::TokenStream) -> Result {
         let mut visitor = AsPrimitiveVisitor::default();
+        let expr: ExprMethodCall = syn::parse2(input).expect("should be a valid method call");
         visitor.visit_expr_method_call(&expr);
         if !visitor.errors.is_empty() {
             return Err(visitor.errors);
@@ -58,21 +59,21 @@ mod tests {
 
     #[test]
     fn test_as_primitives() {
-        let input = r#"x.as_u32()"#;
+        let input = quote! {x.as_u32() };
         assert!(lint(input).is_err());
-        let input = r#"x.as_u64()"#;
+        let input = quote! {x.as_u64() };
         assert!(lint(input).is_err());
-        let input = r#"x.as_u128()"#;
+        let input = quote! {x.as_u128() };
         assert!(lint(input).is_err());
-        let input = r#"x.as_usize()"#;
+        let input = quote! {x.as_usize() };
         assert!(lint(input).is_err());
     }
 
     #[test]
     fn test_non_as_primitives() {
-        let input = r#"x.as_ref()"#;
+        let input = quote! {x.as_ref() };
         assert!(lint(input).is_ok());
-        let input = r#"x.as_slice()"#;
+        let input = quote! {x.as_slice() };
         assert!(lint(input).is_ok());
     }
 }
diff --git a/support/linting/src/forbid_keys_remove.rs b/support/linting/src/forbid_keys_remove.rs
@@ -58,11 +58,12 @@ fn is_keys_remove_call(func: &Expr, args: &Punctuated<Expr, Comma>) -> bool {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use quote::quote;
 
-    fn lint(input: &str) -> Result {
+    fn lint(input: proc_macro2::TokenStream) -> Result {
         let mut visitor = KeysRemoveVisitor::default();
-        visitor
-            .visit_expr_call(&syn::parse_str(input).expect("should only use on a function call"));
+        let expr: syn::ExprCall = syn::parse2(input).expect("should be a valid function call");
+        visitor.visit_expr_call(&expr);
 
         if visitor.errors.is_empty() {
             Ok(())
@@ -73,46 +74,46 @@ mod tests {
 
     #[test]
     fn test_keys_remove_forbidden() {
-        let input = r#"Keys::<T>::remove(netuid, uid_to_replace)"#;
+        let input = quote! { Keys::<T>::remove(netuid, uid_to_replace) };
         assert!(lint(input).is_err());
-        let input = r#"Keys::<U>::remove(netuid, uid_to_replace)"#;
+        let input = quote! { Keys::<U>::remove(netuid, uid_to_replace) };
         assert!(lint(input).is_err());
-        let input = r#"Keys::<U>::remove(1, "2".parse().unwrap(),)"#;
+        let input = quote! { Keys::<U>::remove(1, "2".parse().unwrap(),) };
         assert!(lint(input).is_err());
     }
 
     #[test]
     fn test_non_keys_remove_not_forbidden() {
-        let input = r#"remove(netuid, uid_to_replace)"#;
+        let input = quote! { remove(netuid, uid_to_replace) };
         assert!(lint(input).is_ok());
-        let input = r#"Keys::remove(netuid, uid_to_replace)"#;
+        let input = quote! { Keys::remove(netuid, uid_to_replace) };
         assert!(lint(input).is_ok());
-        let input = r#"Keys::<T>::remove::<U>(netuid, uid_to_replace)"#;
+        let input = quote! { Keys::<T>::remove::<U>(netuid, uid_to_replace) };
         assert!(lint(input).is_ok());
-        let input = r#"Keys::<T>::remove(netuid, uid_to_replace, third_wheel)"#;
+        let input = quote! { Keys::<T>::remove(netuid, uid_to_replace, third_wheel) };
         assert!(lint(input).is_ok());
-        let input = r#"ParentKeys::remove(netuid, uid_to_replace)"#;
+        let input = quote! { ParentKeys::remove(netuid, uid_to_replace) };
         assert!(lint(input).is_ok());
-        let input = r#"ChildKeys::<T>::remove(netuid, uid_to_replace)"#;
+        let input = quote! { ChildKeys::<T>::remove(netuid, uid_to_replace) };
         assert!(lint(input).is_ok());
     }
 
     #[test]
     fn test_keys_remove_allowed() {
-        let input = r#"
-        #[allow(unknown_lints)]
-        Keys::<T>::remove(netuid, uid_to_replace)
-        "#;
+        let input = quote! {
+            #[allow(unknown_lints)]
+            Keys::<T>::remove(netuid, uid_to_replace)
+        };
         assert!(lint(input).is_ok());
-        let input = r#"
-        #[allow(unknown_lints)]
-        Keys::<U>::remove(netuid, uid_to_replace)
-        "#;
+        let input = quote! {
+            #[allow(unknown_lints)]
+            Keys::<U>::remove(netuid, uid_to_replace)
+        };
         assert!(lint(input).is_ok());
-        let input = r#"
-        #[allow(unknown_lints)]
-        Keys::<U>::remove(1, "2".parse().unwrap(),)
-        "#;
+        let input = quote! {
+            #[allow(unknown_lints)]
+            Keys::<U>::remove(1, "2".parse().unwrap(),)
+        };
         assert!(lint(input).is_ok());
     }
 }