extend searching by file name

src/commands/mod.rs

@@ -3,19 +3,19 @@ use structopt::StructOpt;
 
 #[derive(StructOpt)]
 pub struct CLI {
-    // Flag to search all classes with that value
+    /// Flag to search all classes with that value
     #[structopt(
         short = "g",
         long = "grep",
         help = "Used to retrieve all classes with that pattern"
     )]
     pub grep: bool,
 
-    // Class name to be fetched
+    /// Class name to be fetched
     #[structopt(help = "Name of the Python class")]
     pub class_name: String,
 
-    // Search directory
+    /// Search path
     #[structopt(parse(from_os_str), default_value = ".", help = "Search directory")]
-    pub dir_path: PathBuf,
+    pub path: PathBuf,
 }

src/joneslib/loader.rs

@@ -22,15 +22,23 @@ static INHERITANCE_PATTERN: &str = r"<Inherit>\s\[(.*?)\]";
 static OUTPUT_PATTERN: &str = r"<Output> (\w+)";
 
 /// Loads all objects from a Python project, given through the python project path.
-pub fn load_python_project(project_path: &PathBuf) -> Option<Vec<(String, String)>> {
+pub fn load_python_project(path: &PathBuf) -> Option<Vec<(String, String)>> {
     let class_name_pattern = Regex::new(CLASS_NAME_PATTERN).unwrap();
     let file_name_pattern = Regex::new(FILE_NAME_PATTERN).unwrap();
 
-    let script_output = match run_python_script(&project_path) {
-        Some(output) => String::from_utf8(output).unwrap(),
-        None => return None,
+    let script_output = if path.is_dir() {
+        match run_python_script(&path) {
+            Some(output) => String::from_utf8(output).unwrap(),
+            None => return None,
+        }
+    } else {
+        match run_python_single_file_script(&path) {
+            Some(output) => String::from_utf8(output).unwrap(),
+            None => return None,
+        }
     };
 
+
     let file_pattern_captures = file_name_pattern.captures_iter(&script_output);
     let found_classes = class_name_pattern
         .captures_iter(&script_output)
@@ -129,6 +137,36 @@ for root, dirs, files in os.walk({:?}):
     }
 }
 
+#[inline]
+fn run_python_single_file_script(file_path: &PathBuf) -> Option<Vec<u8>> {
+    let python_script = format!(
+        r#"import ast
+import os
+
+file_name = {:?}
+with open(file_name, "r") as file:
+    tree = ast.parse(file.read())
+
+for node in ast.walk(tree):
+    if not isinstance(node, ast.ClassDef):
+        continue
+
+    class_name = node.name
+    print("<Class> %s, <File> %s" % (class_name, file_name))
+"#,
+        file_path.as_os_str(),
+    );
+
+    let output = Command::new("python").arg("-c").arg(python_script).output();
+
+    if let Ok(output) = output {
+        println!("{}", String::from_utf8(output.stderr).unwrap());
+        Some(output.stdout)
+    } else {
+        None
+    }
+}
+
 #[inline]
 fn run_python_class_script(file_path: &PathBuf, class_name: &String) -> Option<Vec<u8>> {
     let python_script = format!(

src/joneslib/mod.rs

@@ -85,8 +85,8 @@ pub fn project_traversal(dir_path: &PathBuf, class_name: &String) -> Option<obje
 
 /// Loads the project classes and filters them by the class name. Returns a vector
 /// of tuples containing the class name and the file path.
-pub fn search(dir_path: &PathBuf, class_name: &String) -> Option<Vec<ClassMatch>> {
-    let project_classes = match loader::load_python_project(dir_path) {
+pub fn search(path: &PathBuf, class_name: &String) -> Option<Vec<ClassMatch>> {
+    let project_classes = match loader::load_python_project(path) {
         Some(classes) => classes,
         None => {
             println!("Error occurred while loading project classes");

src/main.rs

@@ -18,13 +18,13 @@ fn main() {
     let comms = commands::CLI::from_args();
     if comms.grep {
         // Search for a keyword in class name
-        match joneslib::search(&comms.dir_path, &comms.class_name) {
+        match joneslib::search(&comms.path, &comms.class_name) {
             Some(matches) => display::class_matches(matches),
             None => display::not_found_message(),
         }
     } else {
         // Generate python class
-        match joneslib::project_traversal(&comms.dir_path, &comms.class_name) {
+        match joneslib::project_traversal(&comms.path, &comms.class_name) {
             Some(class) => joneslib::display::output_class(&class),
             None => display::not_found_message(),
         }