From 94d2c20dcc8072a5ce0d92d2be2c69e4f1e7677e Mon Sep 17 00:00:00 2001 From: Peter Byfield Date: Mon, 27 Oct 2025 09:42:11 +0100 Subject: [PATCH] Add support for dot format --- README.md | 28 ++++- src/graph.rs | 110 +++++++++++++++++- src/lib.rs | 227 +++++++++++++++++++++++++++++--------- src/main.rs | 101 ++++++++++++++++- tests/integration_test.rs | 132 ++++++++++++++++++++++ 5 files changed, 533 insertions(+), 65 deletions(-) diff --git a/README.md b/README.md index 2aa096f..72ff4a2 100644 --- a/README.md +++ b/README.md @@ -14,7 +14,7 @@ A Rust CLI tool and library for finding all subclasses (direct and transitive) o - **Re-export support**: Tracks classes defined in one module but exported from another - **Ambiguity detection**: Detects when a class name appears in multiple modules and provides clear guidance - **Gitignore support**: Automatically respects `.gitignore` files using the `ignore` crate -- **Multiple output formats**: Text and JSON output formats +- **Multiple output formats**: Text, JSON, and Graphviz dot formats for visualization - **Fast and efficient**: Written in Rust with parallel file traversal - **Smart caching**: Caches parsed files with gzip compression for 2.5x speedup on repeated runs - **Configurable logging**: Uses `env_logger` for flexible logging control @@ -78,6 +78,28 @@ Get results in JSON format for scripting: pysubclasses Animal --format json ``` +### Graphviz Dot Format + +Generate a graph visualization in Graphviz dot format: + +```bash +# Output dot format +pysubclasses Animal --format dot + +# Pipe to dot to generate an image +pysubclasses Animal --format dot | dot -Tpng > graph.png +pysubclasses Animal --format dot | dot -Tsvg > graph.svg + +# Works with both modes +pysubclasses Animal --format dot --mode direct # Shows only direct relationships +pysubclasses Animal --format dot --mode all # Shows full inheritance tree +``` + +The dot format includes: +- The base class (highlighted in green) +- All subclasses (in blue) +- Arrows showing inheritance relationships + ### Search Mode Control whether to find only direct subclasses or all transitive subclasses: @@ -231,12 +253,12 @@ $ pysubclasses Animal --format json You can also use `pysubclasses` as a library in your Rust projects: ```rust -use pysubclasses::{SubclassFinder, SubclassMode}; +use pysubclasses::{SubclassFinder, SearchMode}; use std::path::PathBuf; fn main() -> Result<(), Box> { let finder = SubclassFinder::new(PathBuf::from("./src"))?; - let subclasses = finder.find_subclasses("BaseClass", None, SubclassMode::All)?; + let subclasses = finder.find_subclasses("BaseClass", None, SearchMode::All)?; for class_ref in subclasses { println!("{} ({})", class_ref.class_name, class_ref.module_path); diff --git a/src/graph.rs b/src/graph.rs index 2b3fd31..f061446 100644 --- a/src/graph.rs +++ b/src/graph.rs @@ -10,12 +10,15 @@ use crate::registry::{ClassId, Registry}; /// An inheritance graph representing parent-child relationships between classes. /// -/// This structure maps each class to the set of classes that directly inherit from it. -/// The graph can be used to efficiently find all descendants of a given class using -/// breadth-first search. +/// This structure maps each class to the set of classes that directly inherit from it, +/// and also maintains the reverse mapping for efficient parent lookup. +/// The graph can be used to efficiently find all descendants or ancestors of a given +/// class using breadth-first search. pub struct InheritanceGraph { /// Maps parent classes to their direct children. pub children: HashMap>, + /// Maps child classes to their direct parents. + parents: HashMap>, } impl InheritanceGraph { @@ -41,22 +44,29 @@ impl InheritanceGraph { /// 3. Skip any base classes that cannot be resolved (e.g., external dependencies) pub fn build(registry: &Registry) -> Self { let mut children: HashMap> = HashMap::new(); + let mut parents: HashMap> = HashMap::new(); - // Build parent → children edges by examining each class's bases + // Build parent → children and child → parents edges by examining each class's bases for (child_id, metadata) in ®istry.classes { for base_name in &metadata.bases { // Resolve the base class reference in this class's module context if let Some(parent_id) = registry.resolve_class(&child_id.module, base_name) { // Add this class as a child of its parent children - .entry(parent_id) + .entry(parent_id.clone()) .or_default() .insert(child_id.clone()); + + // Add the parent as a parent of this class + parents + .entry(child_id.clone()) + .or_default() + .insert(parent_id); } } } - Self { children } + Self { children, parents } } /// Finds only the direct subclasses of a given class. @@ -149,4 +159,92 @@ impl InheritanceGraph { result } + + /// Finds only the direct parent classes of a given class. + /// + /// Returns classes that the specified class directly inherits from, + /// without traversing further up the inheritance hierarchy. + /// + /// # Arguments + /// + /// * `root` - The class to find direct parent classes for + /// + /// # Returns + /// + /// A vector containing only the direct parent classes. + /// + /// # Examples + /// + /// ```text + /// Given: + /// class Animal: pass + /// class Mammal(Animal): pass + /// class Dog(Mammal): pass + /// + /// find_direct_parent_classes(Dog) → [Mammal] + /// find_direct_parent_classes(Mammal) → [Animal] + /// ``` + pub fn find_direct_parent_classes(&self, root: &ClassId) -> Vec { + self.parents + .get(root) + .map(|parents| parents.iter().cloned().collect()) + .unwrap_or_default() + } + + /// Finds all transitive parent classes of a given class. + /// + /// Performs a breadth-first search to discover all classes that the specified + /// class directly or indirectly inherits from. This includes: + /// - Direct parents (one level of inheritance) + /// - Grandparents (two levels) + /// - Great-grandparents, etc. (any depth) + /// + /// # Arguments + /// + /// * `root` - The class to find parent classes for + /// + /// # Returns + /// + /// A vector containing all transitive parent classes. The root class itself is not + /// included in the result. The order of classes in the vector is determined by + /// the BFS traversal order. + /// + /// # Examples + /// + /// ```text + /// Given: + /// class Animal: pass + /// class Mammal(Animal): pass + /// class Dog(Mammal): pass + /// + /// find_all_parent_classes(Dog) → [Mammal, Animal] + /// ``` + pub fn find_all_parent_classes(&self, root: &ClassId) -> Vec { + use std::collections::VecDeque; + + let mut result = Vec::new(); + let mut visited = HashSet::new(); + let mut queue = VecDeque::new(); + + // Initialize BFS with the root class + queue.push_back(root.clone()); + visited.insert((root.module.clone(), root.name.clone())); + + // BFS traversal + while let Some(current) = queue.pop_front() { + // Examine all direct parents of the current class + if let Some(parents) = self.parents.get(¤t) { + for parent in parents { + let key = (parent.module.clone(), parent.name.clone()); + if !visited.contains(&key) { + visited.insert(key); + result.push(parent.clone()); + queue.push_back(parent.clone()); + } + } + } + } + + result + } } diff --git a/src/lib.rs b/src/lib.rs index af93a40..9a41145 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,12 +7,12 @@ //! # Examples //! //! ```no_run -//! use pysubclasses::{SubclassFinder, SubclassMode}; +//! use pysubclasses::{SubclassFinder, SearchMode}; //! use std::path::PathBuf; //! //! # fn main() -> Result<(), Box> { //! let finder = SubclassFinder::new(PathBuf::from("./src"))?; -//! let subclasses = finder.find_subclasses("BaseClass", None, SubclassMode::All)?; +//! let subclasses = finder.find_subclasses("BaseClass", None, SearchMode::All)?; //! //! for class_ref in subclasses { //! println!("{} ({})", class_ref.class_name, class_ref.module_path); @@ -53,12 +53,12 @@ impl ClassReference { } } -/// Mode for finding subclasses. +/// Mode for searching the inheritance graph. #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum SubclassMode { - /// Find only direct subclasses (one level of inheritance) +pub enum SearchMode { + /// Find only direct relationships (one level of inheritance) Direct, - /// Find all transitive subclasses (any depth) + /// Find all transitive relationships (any depth) All, } @@ -67,7 +67,7 @@ pub enum SubclassMode { /// # Examples /// /// ```no_run -/// use pysubclasses::{SubclassFinder, SubclassMode}; +/// use pysubclasses::{SubclassFinder, SearchMode}; /// use std::path::PathBuf; /// /// # fn main() -> Result<(), Box> { @@ -75,10 +75,10 @@ pub enum SubclassMode { /// let finder = SubclassFinder::new(PathBuf::from("."))?; /// /// // Find all subclasses of "Animal" -/// let subclasses = finder.find_subclasses("Animal", None, SubclassMode::All)?; +/// let subclasses = finder.find_subclasses("Animal", None, SearchMode::All)?; /// /// // Find direct subclasses of "Animal" from a specific module -/// let subclasses = finder.find_subclasses("Animal", Some("zoo.animals"), SubclassMode::Direct)?; +/// let subclasses = finder.find_subclasses("Animal", Some("zoo.animals"), SearchMode::Direct)?; /// # Ok(()) /// # } /// ``` @@ -179,17 +179,17 @@ impl SubclassFinder { /// # Examples /// /// ```no_run - /// use pysubclasses::{SubclassFinder, SubclassMode}; + /// use pysubclasses::{SubclassFinder, SearchMode}; /// use std::path::PathBuf; /// /// # fn main() -> Result<(), Box> { /// let finder = SubclassFinder::new(PathBuf::from("."))?; /// /// // Find all transitive subclasses - /// let subclasses = finder.find_subclasses("Animal", None, SubclassMode::All)?; + /// let subclasses = finder.find_subclasses("Animal", None, SearchMode::All)?; /// /// // Find only direct subclasses - /// let direct = finder.find_subclasses("Animal", None, SubclassMode::Direct)?; + /// let direct = finder.find_subclasses("Animal", None, SearchMode::Direct)?; /// # Ok(()) /// # } /// ``` @@ -197,51 +197,15 @@ impl SubclassFinder { &self, class_name: &str, module_path: Option<&str>, - mode: SubclassMode, + mode: SearchMode, ) -> Result> { // Find the target class - let target_id = if let Some(module) = module_path { - // Module specified - look for class in that module (including re-exports) - if let Some(resolved_id) = self.registry.resolve_class(module, class_name) { - resolved_id - } else { - return Err(Error::ClassNotFound { - name: class_name.to_string(), - module_path: Some(module.to_string()), - }); - } - } else { - // No module specified - search for class by name - let mut matches = Vec::new(); - for class_id in self.registry.classes.keys() { - if class_id.name == class_name { - matches.push(class_id.clone()); - } - } - - match matches.len() { - 0 => { - return Err(Error::ClassNotFound { - name: class_name.to_string(), - module_path: None, - }); - } - 1 => matches.into_iter().next().unwrap(), - _ => { - let candidates: Vec = - matches.iter().map(|id| id.module.clone()).collect(); - return Err(Error::AmbiguousClassName { - name: class_name.to_string(), - candidates, - }); - } - } - }; + let target_id = self.resolve_target_class(class_name, module_path)?; // Find subclasses using the graph based on the mode let subclass_ids = match mode { - SubclassMode::Direct => self.graph.find_direct_subclasses(&target_id), - SubclassMode::All => self.graph.find_all_subclasses(&target_id), + SearchMode::Direct => self.graph.find_direct_subclasses(&target_id), + SearchMode::All => self.graph.find_all_subclasses(&target_id), }; // Convert to ClassReference @@ -269,8 +233,167 @@ impl SubclassFinder { Ok(results) } + /// Finds parent classes of a given class with a specified mode. + /// + /// # Arguments + /// + /// * `class_name` - The simple name of the class to find parent classes for + /// * `module_path` - Optional module path to disambiguate the class if the name + /// appears multiple times in the codebase + /// * `mode` - Whether to find only direct parents or all transitive parents + /// + /// # Returns + /// + /// A sorted vector of parent classes. The results are sorted by module path for + /// consistent output. + /// + /// # Errors + /// + /// Returns an error if: + /// - The class is not found + /// - The class name is ambiguous and no module path is provided + /// + /// # Examples + /// + /// ```no_run + /// use pysubclasses::{SubclassFinder, SearchMode}; + /// use std::path::PathBuf; + /// + /// # fn main() -> Result<(), Box> { + /// let finder = SubclassFinder::new(PathBuf::from("."))?; + /// + /// // Find all transitive parent classes + /// let parents = finder.find_parent_classes("Dog", Some("animals"), SearchMode::All)?; + /// + /// // Find only direct parent classes + /// let direct = finder.find_parent_classes("Dog", Some("animals"), SearchMode::Direct)?; + /// # Ok(()) + /// # } + /// ``` + pub fn find_parent_classes( + &self, + class_name: &str, + module_path: Option<&str>, + mode: SearchMode, + ) -> Result> { + // Find the target class + let class_id = self.resolve_target_class(class_name, module_path)?; + + // Find parent classes using the graph based on the mode + let parent_ids = match mode { + SearchMode::Direct => self.graph.find_direct_parent_classes(&class_id), + SearchMode::All => self.graph.find_all_parent_classes(&class_id), + }; + + // Convert to ClassReference + let mut results: Vec = parent_ids + .into_iter() + .filter_map(|id| { + self.registry + .modules + .get(&id.module) + .map(|metadata| ClassReference { + class_name: id.name.clone(), + module_path: id.module.clone(), + file_path: metadata.file_path.clone(), + }) + }) + .collect(); + + // Sort by module path for consistent output + results.sort_by(|a, b| { + a.module_path + .cmp(&b.module_path) + .then(a.class_name.cmp(&b.class_name)) + }); + + Ok(results) + } + /// Returns the number of classes found in the codebase. pub fn class_count(&self) -> usize { self.registry.classes.len() } + + /// Resolves a target class by name and optional module path. + /// + /// This helper method encapsulates the logic for finding a class given its name + /// and optional module path. It handles: + /// - Module-qualified lookups (with re-export resolution) + /// - Unqualified lookups by class name + /// - Ambiguity detection when multiple classes have the same name + /// + /// # Arguments + /// + /// * `class_name` - The simple name of the class to find + /// * `module_path` - Optional module path to disambiguate the class + /// + /// # Returns + /// + /// The ClassId of the resolved class. + /// + /// # Errors + /// + /// Returns an error if: + /// - The class is not found + /// - The class name is ambiguous and no module path is provided + fn resolve_target_class( + &self, + class_name: &str, + module_path: Option<&str>, + ) -> Result { + if let Some(module) = module_path { + // Module specified - look for class in that module (including re-exports) + if let Some(resolved_id) = self.registry.resolve_class(module, class_name) { + Ok(resolved_id) + } else { + Err(Error::ClassNotFound { + name: class_name.to_string(), + module_path: Some(module.to_string()), + }) + } + } else { + // No module specified - search for class by name + let mut matches = Vec::new(); + for class_id in self.registry.classes.keys() { + if class_id.name == class_name { + matches.push(class_id.clone()); + } + } + + match matches.len() { + 0 => Err(Error::ClassNotFound { + name: class_name.to_string(), + module_path: None, + }), + 1 => Ok(matches.into_iter().next().unwrap()), + _ => { + let candidates: Vec = + matches.iter().map(|id| id.module.clone()).collect(); + Err(Error::AmbiguousClassName { + name: class_name.to_string(), + candidates, + }) + } + } + } + } + + /// Resolves a class reference by name and optional module. + pub fn resolve_class_reference( + &self, + class_name: &str, + module_path: Option<&str>, + ) -> Option { + // Use resolve_target_class for the core logic + let target_id = self.resolve_target_class(class_name, module_path).ok()?; + + // Convert to ClassReference + let metadata = self.registry.modules.get(&target_id.module)?; + Some(ClassReference { + class_name: target_id.name.clone(), + module_path: target_id.module.clone(), + file_path: metadata.file_path.clone(), + }) + } } diff --git a/src/main.rs b/src/main.rs index 2431170..52858a0 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,7 +2,7 @@ use anyhow::{Context, Result}; use clap::Parser; -use pysubclasses::{ClassReference, SubclassFinder, SubclassMode}; +use pysubclasses::{ClassReference, SearchMode, SubclassFinder}; use serde::Serialize; use std::path::PathBuf; @@ -58,6 +58,8 @@ enum OutputFormat { Text, /// JSON output Json, + /// Graphviz dot format + Dot, } #[derive(Debug, Clone, Copy, clap::ValueEnum)] @@ -116,10 +118,10 @@ fn main() -> Result<()> { .unwrap_or_default() ); - // Convert CLI Mode to SubclassMode + // Convert CLI Mode to SearchMode let mode = match args.mode { - Mode::Direct => SubclassMode::Direct, - Mode::All => SubclassMode::All, + Mode::Direct => SearchMode::Direct, + Mode::All => SearchMode::All, }; // Find subclasses @@ -144,6 +146,7 @@ fn main() -> Result<()> { match args.format { OutputFormat::Text => output_text(&args.class_name, &subclasses), OutputFormat::Json => output_json(&args.class_name, &args.module, &subclasses)?, + OutputFormat::Dot => output_dot(&args.class_name, &args.module, &subclasses, &finder)?, } Ok(()) @@ -187,3 +190,93 @@ fn output_json( println!("{}", serde_json::to_string_pretty(&output)?); Ok(()) } + +fn output_dot( + class_name: &str, + module_path: &Option, + subclasses: &[ClassReference], + finder: &SubclassFinder, +) -> Result<()> { + use std::collections::HashSet; + + // Resolve the actual base class + let base_class = finder + .resolve_class_reference(class_name, module_path.as_deref()) + .with_context(|| format!("Failed to resolve base class '{}'", class_name))?; + + // Create a set of all classes in our graph for quick lookup + let mut class_set = HashSet::new(); + class_set.insert(( + base_class.class_name.clone(), + base_class.module_path.clone(), + )); + for subclass in subclasses { + class_set.insert((subclass.class_name.clone(), subclass.module_path.clone())); + } + + println!("digraph {{"); + println!(" rankdir=TB;"); + println!(" node [shape=box, style=filled, fillcolor=lightblue];"); + println!(); + + // Add base class node + let base_node_id = format!( + "{}_{}", + sanitize_for_dot(&base_class.module_path), + sanitize_for_dot(&base_class.class_name) + ); + println!( + " {} [label=\"{}\\n({})\", fillcolor=lightgreen];", + base_node_id, base_class.class_name, base_class.module_path + ); + + // Add subclass nodes + for subclass in subclasses { + let node_id = format!( + "{}_{}", + sanitize_for_dot(&subclass.module_path), + sanitize_for_dot(&subclass.class_name) + ); + println!( + " {} [label=\"{}\\n({})\"];", + node_id, subclass.class_name, subclass.module_path + ); + } + + println!(); + + // Add edges + for subclass in subclasses { + let parents = finder + .find_parent_classes( + &subclass.class_name, + Some(&subclass.module_path), + SearchMode::Direct, + ) + .unwrap_or_default(); + + for parent in parents { + // Only draw edge if parent is in our class set (either base or another subclass) + if class_set.contains(&(parent.class_name.clone(), parent.module_path.clone())) { + let parent_node_id = format!( + "{}_{}", + sanitize_for_dot(&parent.module_path), + sanitize_for_dot(&parent.class_name) + ); + let child_node_id = format!( + "{}_{}", + sanitize_for_dot(&subclass.module_path), + sanitize_for_dot(&subclass.class_name) + ); + println!(" {} -> {};", parent_node_id, child_node_id); + } + } + } + + println!("}}"); + Ok(()) +} + +fn sanitize_for_dot(s: &str) -> String { + s.replace(['.', '-'], "_") +} diff --git a/tests/integration_test.rs b/tests/integration_test.rs index d44f4db..d49e9de 100644 --- a/tests/integration_test.rs +++ b/tests/integration_test.rs @@ -820,3 +820,135 @@ class Cat(Mammal): temp.close().unwrap(); } + +#[test] +fn test_find_parent_classes_direct() { + use pysubclasses::{SearchMode, SubclassFinder}; + + let temp = assert_fs::TempDir::new().unwrap(); + + // Create inheritance hierarchy: Animal -> Mammal -> Dog + temp.child("animals.py") + .write_str( + r#" +class Animal: + pass + +class Mammal(Animal): + pass + +class Dog(Mammal): + pass +"#, + ) + .unwrap(); + + let finder = SubclassFinder::new(temp.path().to_path_buf()).unwrap(); + + // Dog's direct parent should be Mammal + let parents = finder + .find_parent_classes("Dog", Some("animals"), SearchMode::Direct) + .unwrap(); + assert_eq!(parents.len(), 1); + assert_eq!(parents[0].class_name, "Mammal"); + assert_eq!(parents[0].module_path, "animals"); + + // Mammal's direct parent should be Animal + let parents = finder + .find_parent_classes("Mammal", Some("animals"), SearchMode::Direct) + .unwrap(); + assert_eq!(parents.len(), 1); + assert_eq!(parents[0].class_name, "Animal"); + assert_eq!(parents[0].module_path, "animals"); + + // Animal should have no parents + let parents = finder + .find_parent_classes("Animal", Some("animals"), SearchMode::Direct) + .unwrap(); + assert_eq!(parents.len(), 0); + + temp.close().unwrap(); +} + +#[test] +fn test_find_parent_classes_all() { + use pysubclasses::{SearchMode, SubclassFinder}; + + let temp = assert_fs::TempDir::new().unwrap(); + + // Create inheritance hierarchy: Animal -> Mammal -> Dog + temp.child("animals.py") + .write_str( + r#" +class Animal: + pass + +class Mammal(Animal): + pass + +class Dog(Mammal): + pass +"#, + ) + .unwrap(); + + let finder = SubclassFinder::new(temp.path().to_path_buf()).unwrap(); + + // Dog's all parents should be Mammal and Animal + let parents = finder + .find_parent_classes("Dog", Some("animals"), SearchMode::All) + .unwrap(); + assert_eq!(parents.len(), 2); + + let parent_names: Vec<_> = parents.iter().map(|p| p.class_name.as_str()).collect(); + assert!(parent_names.contains(&"Mammal")); + assert!(parent_names.contains(&"Animal")); + assert!(parents.iter().all(|p| p.module_path == "animals")); + + // Mammal's all parents should be just Animal + let parents = finder + .find_parent_classes("Mammal", Some("animals"), SearchMode::All) + .unwrap(); + assert_eq!(parents.len(), 1); + assert_eq!(parents[0].class_name, "Animal"); + assert_eq!(parents[0].module_path, "animals"); + + temp.close().unwrap(); +} + +#[test] +fn test_find_parent_classes_multiple_inheritance() { + use pysubclasses::{SearchMode, SubclassFinder}; + + let temp = assert_fs::TempDir::new().unwrap(); + + // Create multiple inheritance + temp.child("animals.py") + .write_str( + r#" +class Animal: + pass + +class Walker: + pass + +class Dog(Animal, Walker): + pass +"#, + ) + .unwrap(); + + let finder = SubclassFinder::new(temp.path().to_path_buf()).unwrap(); + + // Dog should have both Animal and Walker as direct parents + let parents = finder + .find_parent_classes("Dog", Some("animals"), SearchMode::Direct) + .unwrap(); + assert_eq!(parents.len(), 2); + + let parent_names: Vec<_> = parents.iter().map(|p| p.class_name.as_str()).collect(); + assert!(parent_names.contains(&"Animal")); + assert!(parent_names.contains(&"Walker")); + + temp.close().unwrap(); +}