codegraph_python/relationships/
imports.rs

1use serde::{Deserialize, Serialize};
2
3/// Represents an import statement
4#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
5pub struct ImportEntity {
6    /// List of imported items (names, modules, etc.)
7    pub imported_items: Vec<String>,
8
9    /// Source module being imported from
10    pub from_module: String,
11
12    /// Line number of the import statement
13    pub line: usize,
14
15    /// Is this a wildcard import (from module import *)?
16    pub is_wildcard: bool,
17}
18
19impl ImportEntity {
20    /// Create a regular import
21    pub fn new(from_module: impl Into<String>, imported_items: Vec<String>, line: usize) -> Self {
22        Self {
23            imported_items,
24            from_module: from_module.into(),
25            line,
26            is_wildcard: false,
27        }
28    }
29
30    /// Create a wildcard import
31    pub fn wildcard(from_module: impl Into<String>, line: usize) -> Self {
32        Self {
33            imported_items: vec![],
34            from_module: from_module.into(),
35            line,
36            is_wildcard: true,
37        }
38    }
39
40    /// Check if this imports a specific name
41    pub fn imports_name(&self, name: &str) -> bool {
42        if self.is_wildcard {
43            true // Wildcard imports everything
44        } else {
45            self.imported_items.iter().any(|item| item == name)
46        }
47    }
48
49    /// Add an imported item
50    pub fn add_item(mut self, item: impl Into<String>) -> Self {
51        self.imported_items.push(item.into());
52        self
53    }
54}
55
56#[cfg(test)]
57mod tests {
58    use super::*;
59
60    #[test]
61    fn test_import_entity_new() {
62        let import = ImportEntity::new("os", vec!["path".to_string()], 1);
63        assert_eq!(import.from_module, "os");
64        assert_eq!(import.imported_items, vec!["path"]);
65        assert_eq!(import.line, 1);
66        assert!(!import.is_wildcard);
67    }
68
69    #[test]
70    fn test_wildcard_import() {
71        let import = ImportEntity::wildcard("os", 1);
72        assert_eq!(import.from_module, "os");
73        assert!(import.imported_items.is_empty());
74        assert!(import.is_wildcard);
75    }
76
77    #[test]
78    fn test_imports_name() {
79        let import = ImportEntity::new("os", vec!["path".to_string(), "environ".to_string()], 1);
80        assert!(import.imports_name("path"));
81        assert!(import.imports_name("environ"));
82        assert!(!import.imports_name("getcwd"));
83
84        let wildcard = ImportEntity::wildcard("os", 1);
85        assert!(wildcard.imports_name("anything"));
86    }
87}
codegraph_python/relationships/imports.rs

codegraph_python/relationships/
imports.rs
