python_ast/
scope.rs

1use std::collections::HashMap;
2use std::default::Default;
3
4/// Represents a single symbol within a scope.
5#[derive(Clone, Debug, Default)]
6pub enum Symbol {
7    SubScope(Box<Scope>),
8    Class(),
9    Variable(),
10    Const(String),
11    #[default]
12    Unknown,
13}
14
15/// Python uses LEGB scope: Local, Enclosing, Global, and Built-in.
16/// Local scope consists of local variables inside a function. Names in the local scope may change new declarations overwrite older ones.
17/// Enclosing scope is the scope of a containing function with inner/nested functions.
18/// Global is the global scope. Names within the global namespace must be unique.
19/// Built-in is basically a special scope for elements that are built into Python.
20#[derive(Clone, Debug, Default)]
21pub enum Scope {
22    #[default]
23    None,
24    Local(HashMap<String, Symbol>),
25    Global(HashMap<String, Symbol>),
26}
27
28impl Scope {
29    pub fn new_local() -> Self {
30        Scope::Local(HashMap::new())
31    }
32
33    pub fn new_global() -> Self {
34        Scope::Global(HashMap::new())
35    }
36
37    pub fn insert(&mut self, key: String, symbol: Symbol) -> Option<Symbol> {
38        match self {
39            Scope::Local(map) | Scope::Global(map) => map.insert(key, symbol),
40            Scope::None => None,
41        }
42    }
43
44    pub fn get(&self, key: &str) -> Option<&Symbol> {
45        match self {
46            Scope::Local(map) | Scope::Global(map) => map.get(key),
47            Scope::None => None,
48        }
49    }
50
51    pub fn contains_key(&self, key: &str) -> bool {
52        match self {
53            Scope::Local(map) | Scope::Global(map) => map.contains_key(key),
54            Scope::None => false,
55        }
56    }
57
58    pub fn len(&self) -> usize {
59        match self {
60            Scope::Local(map) | Scope::Global(map) => map.len(),
61            Scope::None => 0,
62        }
63    }
64
65    pub fn is_empty(&self) -> bool {
66        self.len() == 0
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73
74    #[test]
75    fn test_symbol_default() {
76        let symbol = Symbol::default();
77        matches!(symbol, Symbol::Unknown);
78    }
79
80    #[test]
81    fn test_symbol_variants() {
82        let class_symbol = Symbol::Class();
83        let variable_symbol = Symbol::Variable();
84        let const_symbol = Symbol::Const("test_value".to_string());
85        let unknown_symbol = Symbol::Unknown;
86
87        matches!(class_symbol, Symbol::Class());
88        matches!(variable_symbol, Symbol::Variable());
89        matches!(const_symbol, Symbol::Const(_));
90        matches!(unknown_symbol, Symbol::Unknown);
91    }
92
93    #[test]
94    fn test_symbol_subscope() {
95        let inner_scope = Box::new(Scope::new_local());
96        let subscope_symbol = Symbol::SubScope(inner_scope);
97        
98        match subscope_symbol {
99            Symbol::SubScope(scope) => {
100                assert!(scope.is_empty());
101            }
102            _ => panic!("Expected SubScope symbol"),
103        }
104    }
105
106    #[test]
107    fn test_scope_default() {
108        let scope = Scope::default();
109        matches!(scope, Scope::None);
110        assert!(scope.is_empty());
111    }
112
113    #[test]
114    fn test_scope_new_local() {
115        let scope = Scope::new_local();
116        matches!(scope, Scope::Local(_));
117        assert!(scope.is_empty());
118    }
119
120    #[test]
121    fn test_scope_new_global() {
122        let scope = Scope::new_global();
123        matches!(scope, Scope::Global(_));
124        assert!(scope.is_empty());
125    }
126
127    #[test]
128    fn test_scope_insert_and_get_local() {
129        let mut scope = Scope::new_local();
130        let symbol = Symbol::Variable();
131        
132        let previous = scope.insert("test_var".to_string(), symbol);
133        assert!(previous.is_none());
134        assert_eq!(scope.len(), 1);
135        
136        let retrieved = scope.get("test_var");
137        assert!(retrieved.is_some());
138        matches!(retrieved.unwrap(), Symbol::Variable());
139    }
140
141    #[test]
142    fn test_scope_insert_and_get_global() {
143        let mut scope = Scope::new_global();
144        let symbol = Symbol::Class();
145        
146        let previous = scope.insert("TestClass".to_string(), symbol);
147        assert!(previous.is_none());
148        assert_eq!(scope.len(), 1);
149        
150        let retrieved = scope.get("TestClass");
151        assert!(retrieved.is_some());
152        matches!(retrieved.unwrap(), Symbol::Class());
153    }
154
155    #[test]
156    fn test_scope_insert_overwrite() {
157        let mut scope = Scope::new_local();
158        let symbol1 = Symbol::Variable();
159        let symbol2 = Symbol::Const("new_value".to_string());
160        
161        scope.insert("test_key".to_string(), symbol1);
162        let previous = scope.insert("test_key".to_string(), symbol2);
163        
164        assert!(previous.is_some());
165        matches!(previous.unwrap(), Symbol::Variable());
166        assert_eq!(scope.len(), 1);
167        
168        let current = scope.get("test_key");
169        match current.unwrap() {
170            Symbol::Const(value) => assert_eq!(value, "new_value"),
171            _ => panic!("Expected Const symbol"),
172        }
173    }
174
175    #[test]
176    fn test_scope_none_operations() {
177        let mut scope = Scope::None;
178        let symbol = Symbol::Variable();
179        
180        let result = scope.insert("test".to_string(), symbol);
181        assert!(result.is_none());
182        
183        let retrieved = scope.get("test");
184        assert!(retrieved.is_none());
185        
186        assert!(!scope.contains_key("test"));
187        assert_eq!(scope.len(), 0);
188        assert!(scope.is_empty());
189    }
190
191    #[test]
192    fn test_scope_contains_key() {
193        let mut scope = Scope::new_local();
194        let symbol = Symbol::Variable();
195        
196        assert!(!scope.contains_key("test_var"));
197        
198        scope.insert("test_var".to_string(), symbol);
199        assert!(scope.contains_key("test_var"));
200        assert!(!scope.contains_key("other_var"));
201    }
202
203    #[test]
204    fn test_scope_len_and_empty() {
205        let mut scope = Scope::new_local();
206        
207        assert_eq!(scope.len(), 0);
208        assert!(scope.is_empty());
209        
210        scope.insert("var1".to_string(), Symbol::Variable());
211        assert_eq!(scope.len(), 1);
212        assert!(!scope.is_empty());
213        
214        scope.insert("var2".to_string(), Symbol::Class());
215        assert_eq!(scope.len(), 2);
216        assert!(!scope.is_empty());
217    }
218
219    #[test]
220    fn test_scope_get_nonexistent() {
221        let scope = Scope::new_local();
222        assert!(scope.get("nonexistent").is_none());
223        
224        let global_scope = Scope::new_global();
225        assert!(global_scope.get("nonexistent").is_none());
226        
227        let none_scope = Scope::None;
228        assert!(none_scope.get("nonexistent").is_none());
229    }
230
231    #[test]
232    fn test_nested_scopes() {
233        let mut inner_scope = Scope::new_local();
234        inner_scope.insert("inner_var".to_string(), Symbol::Variable());
235        
236        let mut outer_scope = Scope::new_global();
237        let subscope_symbol = Symbol::SubScope(Box::new(inner_scope));
238        outer_scope.insert("inner_function".to_string(), subscope_symbol);
239        
240        assert_eq!(outer_scope.len(), 1);
241        
242        match outer_scope.get("inner_function").unwrap() {
243            Symbol::SubScope(scope) => {
244                assert_eq!(scope.len(), 1);
245                assert!(scope.contains_key("inner_var"));
246            }
247            _ => panic!("Expected SubScope symbol"),
248        }
249    }
250}