fn visitor update

dcodesdev · Jun 10, 2024 · dd20167 · dd20167
1 parent 339731b
commit dd20167
Showing 1 changed file with 14 additions and 8 deletions.
diff --git a/crates/syntest/src/visitor.rs b/crates/syntest/src/visitor.rs
@@ -17,13 +17,11 @@ impl AsVisitor {
             .find_map(|item| {
                 if let syn::Item::Fn(item_fn) = item {
                     if item_fn.sig.ident.to_string() == fn_name {
-                        Some(item_fn)
-                    } else {
-                        None
+                        return Some(item_fn);
                     }
-                } else {
-                    None
                 }
+
+                None
             })
             .expect(&format!("Function {fn_name} not found"));
 
@@ -56,10 +54,18 @@ pub struct FnVisitor {
 }
 
 impl FnVisitor {
-    pub fn new() -> Self {
-        Self {
+    pub fn new(file: &syn::File) -> Self {
+        let mut visitor = Self {
             functions: Vec::new(),
-        }
+        };
+
+        visitor.visit_file(file);
+
+        visitor
+    }
+
+    pub fn functions(&self) -> &Vec<String> {
+        &self.functions
     }
 }