1use std::path::{Path, PathBuf};
2
3use tree_sitter::Language;
4
5use astmap_core::model::SymbolKind;
6
7use super::{
8 find_enclosing_symbol, first_line, node_text, parse_with_tree_sitter, symbol_from_node,
9 ExtractedCall, ExtractedImport, ExtractedSymbol, ExtractedTypeRef, LanguageParser,
10 LanguageResolver, LanguageSupport, ParseResult,
11};
12
13pub(super) fn lang() -> LanguageSupport {
14 LanguageSupport {
15 name: "python",
16 extensions: &["py"],
17 parser: &PythonParser,
18 resolver_factory: |_root| Box::new(PythonResolver),
19 config_files: &["pyproject.toml"],
20 sibling_fn: python_sibling_expansion,
21 }
22}
23
24fn python_sibling_expansion(rel_path: &str) -> Option<astmap_core::SiblingExpansion> {
25 let dir = rel_path
26 .strip_suffix("__init__.py")
27 .filter(|p| !p.is_empty())?;
28 Some(astmap_core::SiblingExpansion {
29 prefix: dir.to_string(),
30 extension: "py".to_string(),
31 root_only: false,
32 })
33}
34
35pub(crate) struct PythonParser;
38
39impl LanguageParser for PythonParser {
40 fn parse(&self, source: &str, _file_path: &Path) -> ParseResult {
41 parse_python_file(source)
42 }
43
44 fn language_name(&self) -> &str {
45 "python"
46 }
47}
48
49fn python_language() -> Language {
50 tree_sitter_python::LANGUAGE.into()
51}
52
53fn parse_python_file(source: &str) -> ParseResult {
54 parse_with_tree_sitter(
55 source,
56 &python_language(),
57 |root, src, symbols, imports| extract_from_node(root, src, symbols, imports, None),
58 extract_calls,
59 extract_type_refs,
60 )
61}
62
63fn extract_from_node(
68 node: tree_sitter::Node,
69 source: &str,
70 symbols: &mut Vec<ExtractedSymbol>,
71 imports: &mut Vec<ExtractedImport>,
72 parent_index: Option<usize>,
73) {
74 match node.kind() {
75 "function_definition" => {
76 if let Some(sym) = extract_function(&node, source, parent_index) {
77 symbols.push(sym);
78 return; }
80 }
81 "class_definition" => {
82 if let Some(sym) = extract_class(&node, source, parent_index) {
83 let idx = symbols.len();
84 symbols.push(sym);
85 if let Some(body) = node.child_by_field_name("body") {
87 for i in 0..body.child_count() {
88 if let Some(child) = body.child(i as u32) {
89 extract_from_node(child, source, symbols, imports, Some(idx));
90 }
91 }
92 }
93 return;
94 }
95 }
96 "decorated_definition" => {
97 extract_decorated(&node, source, symbols, imports, parent_index);
98 return;
99 }
100 "expression_statement" if parent_index.is_none() => {
101 for i in 0..node.child_count() {
103 if let Some(child) = node.child(i as u32) {
104 if child.kind() == "assignment" {
105 if let Some(sym) = extract_module_constant(&child, &node, source) {
106 symbols.push(sym);
107 }
108 }
109 }
110 }
111 }
112 "import_statement" => {
113 imports.extend(extract_import_plain(&node, source));
114 }
115 "import_from_statement" => {
116 if let Some(imp) = extract_import_from(&node, source) {
117 imports.push(imp);
118 }
119 }
120 _ => {}
121 }
122
123 for i in 0..node.child_count() {
125 if let Some(child) = node.child(i as u32) {
126 extract_from_node(child, source, symbols, imports, parent_index);
127 }
128 }
129}
130
131fn extract_function(
132 node: &tree_sitter::Node,
133 source: &str,
134 parent_index: Option<usize>,
135) -> Option<ExtractedSymbol> {
136 let name_node = node.child_by_field_name("name")?;
137 let name = node_text(&name_node, source);
138 let kind = if parent_index.is_some() {
139 SymbolKind::Method
140 } else {
141 SymbolKind::Function
142 };
143 Some(symbol_from_node(name, kind, node, source, parent_index))
144}
145
146fn extract_class(
147 node: &tree_sitter::Node,
148 source: &str,
149 parent_index: Option<usize>,
150) -> Option<ExtractedSymbol> {
151 let name_node = node.child_by_field_name("name")?;
152 let name = node_text(&name_node, source);
153 Some(symbol_from_node(
154 name,
155 SymbolKind::Class,
156 node,
157 source,
158 parent_index,
159 ))
160}
161
162fn extract_decorated(
163 node: &tree_sitter::Node,
164 source: &str,
165 symbols: &mut Vec<ExtractedSymbol>,
166 imports: &mut Vec<ExtractedImport>,
167 parent_index: Option<usize>,
168) {
169 let decorators: Vec<String> = (0..node.child_count())
171 .filter_map(|i| node.child(i as u32))
172 .filter(|c| c.kind() == "decorator")
173 .map(|c| node_text(&c, source).trim().to_string())
174 .collect();
175
176 let def_node = match node.child_by_field_name("definition") {
177 Some(d) => d,
178 None => return,
179 };
180
181 let decorator_prefix = if decorators.is_empty() {
182 String::new()
183 } else {
184 format!("{}\n", decorators.join("\n"))
185 };
186
187 match def_node.kind() {
188 "function_definition" => {
189 if let Some(mut sym) = extract_function(&def_node, source, parent_index) {
190 sym.signature = Some(format!(
191 "{}{}",
192 decorator_prefix,
193 first_line(&def_node, source)
194 ));
195 sym.start_line = node.start_position().row + 1;
196 sym.start_byte = node.start_byte();
197 symbols.push(sym);
198 }
199 }
200 "class_definition" => {
201 if let Some(mut sym) = extract_class(&def_node, source, parent_index) {
202 sym.signature = Some(format!(
203 "{}{}",
204 decorator_prefix,
205 first_line(&def_node, source)
206 ));
207 sym.start_line = node.start_position().row + 1;
208 sym.start_byte = node.start_byte();
209 let idx = symbols.len();
210 symbols.push(sym);
211 if let Some(body) = def_node.child_by_field_name("body") {
212 for i in 0..body.child_count() {
213 if let Some(child) = body.child(i as u32) {
214 extract_from_node(child, source, symbols, imports, Some(idx));
215 }
216 }
217 }
218 }
219 }
220 _ => {}
221 }
222}
223
224fn extract_module_constant(
225 assignment: &tree_sitter::Node,
226 stmt: &tree_sitter::Node,
227 source: &str,
228) -> Option<ExtractedSymbol> {
229 let left = assignment.child_by_field_name("left")?;
230 if left.kind() != "identifier" {
231 return None;
232 }
233 let name = node_text(&left, source);
234
235 let is_all_caps = name.chars().any(|c| c.is_ascii_uppercase())
236 && name.chars().all(|c| c.is_ascii_uppercase() || c == '_');
237 let has_type_annotation = assignment.child_by_field_name("type").is_some();
238
239 if !is_all_caps && !has_type_annotation {
240 return None;
241 }
242
243 Some(symbol_from_node(
244 name,
245 SymbolKind::Variable,
246 stmt,
247 source,
248 None,
249 ))
250}
251
252fn extract_import_plain(node: &tree_sitter::Node, source: &str) -> Vec<ExtractedImport> {
257 let mut imports = Vec::new();
258 for i in 0..node.child_count() {
259 let child = match node.child(i as u32) {
260 Some(c) => c,
261 None => continue,
262 };
263 match child.kind() {
264 "dotted_name" => {
265 let full_path = node_text(&child, source);
266 let last_segment = full_path
267 .rsplit('.')
268 .next()
269 .unwrap_or(&full_path)
270 .to_string();
271 imports.push(ExtractedImport {
272 path: full_path,
273 imported_symbols: vec![last_segment],
274 });
275 }
276 "aliased_import" => {
277 let name_node = child.child_by_field_name("name");
278 let alias_node = child.child_by_field_name("alias");
279 if let Some(name) = name_node {
280 let path = node_text(&name, source);
281 let symbol = alias_node
282 .map(|a| node_text(&a, source))
283 .unwrap_or_else(|| path.rsplit('.').next().unwrap_or(&path).to_string());
284 imports.push(ExtractedImport {
285 path,
286 imported_symbols: vec![symbol],
287 });
288 }
289 }
290 _ => {}
291 }
292 }
293 imports
294}
295
296fn extract_import_from(node: &tree_sitter::Node, source: &str) -> Option<ExtractedImport> {
297 let module_node = node.child_by_field_name("module_name")?;
298 let path = node_text(&module_node, source);
299
300 let has_wildcard = (0..node.child_count())
302 .filter_map(|i| node.child(i as u32))
303 .any(|c| c.kind() == "wildcard_import");
304
305 if has_wildcard {
306 return Some(ExtractedImport {
307 path,
308 imported_symbols: vec![],
309 });
310 }
311
312 let module_node_id = module_node.id();
314 let mut imported_symbols = Vec::new();
315 for i in 0..node.child_count() {
316 let child = match node.child(i as u32) {
317 Some(c) => c,
318 None => continue,
319 };
320 match child.kind() {
321 "dotted_name" if child.id() != module_node_id => {
322 imported_symbols.push(node_text(&child, source));
323 }
324 "aliased_import" => {
325 if let Some(alias) = child.child_by_field_name("alias") {
326 imported_symbols.push(node_text(&alias, source));
327 } else if let Some(name) = child.child_by_field_name("name") {
328 let text = node_text(&name, source);
329 let last = text.rsplit('.').next().unwrap_or(&text).to_string();
330 imported_symbols.push(last);
331 }
332 }
333 _ => {}
334 }
335 }
336
337 Some(ExtractedImport {
338 path,
339 imported_symbols,
340 })
341}
342
343fn extract_calls(
348 node: tree_sitter::Node,
349 source: &str,
350 symbols: &[ExtractedSymbol],
351) -> Vec<ExtractedCall> {
352 let mut calls = Vec::new();
353 collect_calls(node, source, symbols, &mut calls);
354 calls
355}
356
357fn collect_calls(
358 node: tree_sitter::Node,
359 source: &str,
360 symbols: &[ExtractedSymbol],
361 calls: &mut Vec<ExtractedCall>,
362) {
363 if node.kind() == "call" {
364 if let Some(func_node) = node.child_by_field_name("function") {
365 let callee_name = match func_node.kind() {
366 "attribute" => func_node
367 .child_by_field_name("attribute")
368 .map(|a| node_text(&a, source))
369 .unwrap_or_else(|| node_text(&func_node, source)),
370 "identifier" => node_text(&func_node, source),
371 _ => node_text(&func_node, source),
372 };
373
374 if !is_python_builtin(&callee_name) {
375 let call_line = node.start_position().row + 1;
376 if let Some(caller) = find_enclosing_symbol(symbols, call_line) {
377 calls.push(ExtractedCall {
378 caller_name: caller.to_string(),
379 callee_name,
380 line: call_line,
381 });
382 }
383 }
384 }
385 }
386
387 for i in 0..node.child_count() {
388 if let Some(child) = node.child(i as u32) {
389 collect_calls(child, source, symbols, calls);
390 }
391 }
392}
393
394fn is_python_builtin(name: &str) -> bool {
395 matches!(
396 name,
397 "print"
398 | "len"
399 | "range"
400 | "enumerate"
401 | "zip"
402 | "map"
403 | "filter"
404 | "sorted"
405 | "reversed"
406 | "list"
407 | "dict"
408 | "set"
409 | "tuple"
410 | "str"
411 | "int"
412 | "float"
413 | "bool"
414 | "bytes"
415 | "bytearray"
416 | "type"
417 | "isinstance"
418 | "issubclass"
419 | "hasattr"
420 | "getattr"
421 | "setattr"
422 | "delattr"
423 | "callable"
424 | "super"
425 | "property"
426 | "staticmethod"
427 | "classmethod"
428 | "abs"
429 | "all"
430 | "any"
431 | "ascii"
432 | "bin"
433 | "chr"
434 | "complex"
435 | "dir"
436 | "divmod"
437 | "format"
438 | "globals"
439 | "hash"
440 | "hex"
441 | "id"
442 | "input"
443 | "iter"
444 | "max"
445 | "min"
446 | "next"
447 | "oct"
448 | "open"
449 | "ord"
450 | "pow"
451 | "repr"
452 | "round"
453 | "slice"
454 | "sum"
455 | "vars"
456 | "object"
457 | "Exception"
458 | "ValueError"
459 | "TypeError"
460 | "KeyError"
461 | "IndexError"
462 | "AttributeError"
463 | "RuntimeError"
464 | "StopIteration"
465 | "NotImplementedError"
466 | "OSError"
467 | "IOError"
468 )
469}
470
471fn extract_type_refs(
476 node: tree_sitter::Node,
477 source: &str,
478 symbols: &[ExtractedSymbol],
479) -> Vec<ExtractedTypeRef> {
480 let mut refs = Vec::new();
481 collect_type_refs(node, source, symbols, &mut refs);
482 refs
483}
484
485fn collect_type_refs(
486 node: tree_sitter::Node,
487 source: &str,
488 symbols: &[ExtractedSymbol],
489 refs: &mut Vec<ExtractedTypeRef>,
490) {
491 if node.kind() == "type" {
492 collect_type_identifiers(node, source, symbols, refs);
493 return; }
495
496 for i in 0..node.child_count() {
497 if let Some(child) = node.child(i as u32) {
498 collect_type_refs(child, source, symbols, refs);
499 }
500 }
501}
502
503fn collect_type_identifiers(
504 node: tree_sitter::Node,
505 source: &str,
506 symbols: &[ExtractedSymbol],
507 refs: &mut Vec<ExtractedTypeRef>,
508) {
509 if node.kind() == "identifier" {
510 let name = node_text(&node, source);
511 if !is_python_builtin_type(&name) {
512 let line = node.start_position().row + 1;
513 if let Some(from_sym) = find_enclosing_symbol(symbols, line) {
514 refs.push(ExtractedTypeRef {
515 from_symbol: from_sym.to_string(),
516 to_type: name,
517 line,
518 });
519 }
520 }
521 }
522
523 for i in 0..node.child_count() {
524 if let Some(child) = node.child(i as u32) {
525 collect_type_identifiers(child, source, symbols, refs);
526 }
527 }
528}
529
530fn is_python_builtin_type(name: &str) -> bool {
531 matches!(
532 name,
533 "int"
534 | "float"
535 | "str"
536 | "bool"
537 | "bytes"
538 | "bytearray"
539 | "None"
540 | "list"
541 | "dict"
542 | "set"
543 | "tuple"
544 | "frozenset"
545 | "object"
546 | "type"
547 | "Any"
548 | "Union"
549 | "Optional"
550 | "List"
551 | "Dict"
552 | "Set"
553 | "Tuple"
554 | "FrozenSet"
555 | "Type"
556 | "Callable"
557 | "Iterator"
558 | "Generator"
559 | "Coroutine"
560 | "Awaitable"
561 | "Sequence"
562 | "Mapping"
563 | "MutableMapping"
564 | "Iterable"
565 | "ClassVar"
566 | "Final"
567 | "Literal"
568 | "Protocol"
569 | "Self"
570 | "Never"
571 | "NoReturn"
572 | "TypeVar"
573 )
574}
575
576pub(crate) struct PythonResolver;
579
580impl LanguageResolver for PythonResolver {
581 fn resolve_import(
582 &self,
583 import_path: &str,
584 source_file: &str,
585 project_root: &Path,
586 ) -> Option<PathBuf> {
587 resolve_python_import(import_path, source_file, project_root)
588 }
589}
590
591fn resolve_python_import(
592 import_path: &str,
593 source_file: &str,
594 project_root: &Path,
595) -> Option<PathBuf> {
596 if import_path.is_empty() {
597 return None;
598 }
599
600 if import_path.starts_with('.') {
601 resolve_relative_import(import_path, source_file, project_root)
602 } else {
603 resolve_absolute_import(import_path, project_root)
604 }
605}
606
607fn resolve_absolute_import(import_path: &str, project_root: &Path) -> Option<PathBuf> {
608 let parts: Vec<&str> = import_path.split('.').collect();
609
610 if let Some(resolved) = resolve_module_path(&parts, project_root) {
612 return Some(resolved);
613 }
614
615 for src_dir in &["src", "lib"] {
617 let base = project_root.join(src_dir);
618 if base.is_dir() {
619 if let Some(resolved) = resolve_module_path(&parts, &base) {
620 return Some(resolved);
621 }
622 }
623 }
624
625 None
627}
628
629fn resolve_relative_import(
630 import_path: &str,
631 source_file: &str,
632 project_root: &Path,
633) -> Option<PathBuf> {
634 let dot_count = import_path.chars().take_while(|&c| c == '.').count();
635 let remainder = &import_path[dot_count..];
636
637 let source_abs = project_root.join(source_file);
638 let source_dir = source_abs.parent()?;
639
640 let mut base = source_dir.to_path_buf();
643 for _ in 1..dot_count {
644 base = base.parent()?.to_path_buf();
645 }
646
647 if remainder.is_empty() {
648 let init = base.join("__init__.py");
650 if init.exists() {
651 return init.canonicalize().ok();
652 }
653 return None;
654 }
655
656 let parts: Vec<&str> = remainder.split('.').collect();
657 resolve_module_path(&parts, &base)
658}
659
660fn resolve_module_path(parts: &[&str], base: &Path) -> Option<PathBuf> {
661 if parts.is_empty() {
662 let init = base.join("__init__.py");
663 if init.exists() {
664 return init.canonicalize().ok();
665 }
666 return None;
667 }
668
669 let mut current = base.to_path_buf();
670
671 for (i, part) in parts.iter().enumerate() {
672 if i == parts.len() - 1 {
673 let as_file = current.join(format!("{}.py", part));
675 if as_file.exists() {
676 return as_file.canonicalize().ok();
677 }
678 let as_package = current.join(part).join("__init__.py");
679 if as_package.exists() {
680 return as_package.canonicalize().ok();
681 }
682 return None;
683 } else {
684 let next = current.join(part);
686 if next.is_dir() {
687 current = next;
688 } else {
689 return None;
690 }
691 }
692 }
693 None
694}
695
696#[cfg(test)]
701mod tests {
702 use super::*;
703
704 fn parse(source: &str) -> ParseResult {
707 PythonParser.parse(source, &PathBuf::from("test.py"))
708 }
709
710 #[test]
713 fn test_parse_function() {
714 let result = parse("def hello():\n pass\n");
715 assert!(!result.has_errors);
716 assert_eq!(result.symbols.len(), 1);
717 assert_eq!(result.symbols[0].name, "hello");
718 assert_eq!(result.symbols[0].kind, SymbolKind::Function);
719 assert!(result.symbols[0].parent_index.is_none());
720 }
721
722 #[test]
723 fn test_parse_async_function() {
724 let result = parse("async def fetch():\n pass\n");
725 assert!(!result.has_errors);
726 assert_eq!(result.symbols.len(), 1);
727 assert_eq!(result.symbols[0].name, "fetch");
728 assert_eq!(result.symbols[0].kind, SymbolKind::Function);
729 assert!(result.symbols[0]
730 .signature
731 .as_ref()
732 .unwrap()
733 .contains("async"));
734 }
735
736 #[test]
737 fn test_parse_class_with_methods() {
738 let result = parse(
739 "class Foo:\n def bar(self):\n pass\n def baz(self):\n pass\n",
740 );
741 assert!(!result.has_errors);
742 assert_eq!(result.symbols.len(), 3);
743 assert_eq!(result.symbols[0].name, "Foo");
744 assert_eq!(result.symbols[0].kind, SymbolKind::Class);
745 assert_eq!(result.symbols[1].name, "bar");
746 assert_eq!(result.symbols[1].kind, SymbolKind::Method);
747 assert_eq!(result.symbols[1].parent_index, Some(0));
748 assert_eq!(result.symbols[2].name, "baz");
749 assert_eq!(result.symbols[2].kind, SymbolKind::Method);
750 assert_eq!(result.symbols[2].parent_index, Some(0));
751 }
752
753 #[test]
754 fn test_parse_class_basic() {
755 let result = parse("class Empty:\n pass\n");
756 assert!(!result.has_errors);
757 assert_eq!(result.symbols.len(), 1);
758 assert_eq!(result.symbols[0].name, "Empty");
759 assert_eq!(result.symbols[0].kind, SymbolKind::Class);
760 }
761
762 #[test]
763 fn test_parse_nested_class() {
764 let result = parse("class Outer:\n class Inner:\n pass\n");
765 assert!(!result.has_errors);
766 assert_eq!(result.symbols.len(), 2);
767 assert_eq!(result.symbols[0].name, "Outer");
768 assert_eq!(result.symbols[1].name, "Inner");
769 assert_eq!(result.symbols[1].kind, SymbolKind::Class);
770 assert_eq!(result.symbols[1].parent_index, Some(0));
771 }
772
773 #[test]
774 fn test_decorated_function() {
775 let result = parse("@app.route(\"/\")\ndef index():\n pass\n");
776 assert!(!result.has_errors);
777 assert_eq!(result.symbols.len(), 1);
778 assert_eq!(result.symbols[0].name, "index");
779 assert_eq!(result.symbols[0].kind, SymbolKind::Function);
780 let sig = result.symbols[0].signature.as_ref().unwrap();
781 assert!(sig.contains("@app.route"), "sig = {}", sig);
782 assert!(sig.contains("def index"), "sig = {}", sig);
783 }
784
785 #[test]
786 fn test_multiple_decorators() {
787 let result = parse("@login_required\n@admin_only\ndef admin_view():\n pass\n");
788 assert!(!result.has_errors);
789 assert_eq!(result.symbols.len(), 1);
790 let sig = result.symbols[0].signature.as_ref().unwrap();
791 assert!(sig.contains("@login_required"), "sig = {}", sig);
792 assert!(sig.contains("@admin_only"), "sig = {}", sig);
793 }
794
795 #[test]
796 fn test_decorated_method_in_class() {
797 let result = parse("class Foo:\n @staticmethod\n def helper():\n pass\n");
798 assert!(!result.has_errors);
799 assert_eq!(result.symbols.len(), 2);
800 assert_eq!(result.symbols[1].name, "helper");
801 assert_eq!(result.symbols[1].kind, SymbolKind::Method);
802 assert_eq!(result.symbols[1].parent_index, Some(0));
803 let sig = result.symbols[1].signature.as_ref().unwrap();
804 assert!(sig.contains("@staticmethod"), "sig = {}", sig);
805 }
806
807 #[test]
808 fn test_module_constant_allcaps() {
809 let result = parse("MAX_SIZE = 100\n");
810 assert!(!result.has_errors);
811 assert_eq!(result.symbols.len(), 1);
812 assert_eq!(result.symbols[0].name, "MAX_SIZE");
813 assert_eq!(result.symbols[0].kind, SymbolKind::Variable);
814 }
815
816 #[test]
817 fn test_module_constant_typed() {
818 let result = parse("config: Config = get_config()\n");
819 assert!(!result.has_errors);
820 assert_eq!(result.symbols.len(), 1);
821 assert_eq!(result.symbols[0].name, "config");
822 assert_eq!(result.symbols[0].kind, SymbolKind::Variable);
823 }
824
825 #[test]
826 fn test_non_constant_skipped() {
827 let result = parse("result = compute()\n");
828 assert!(!result.has_errors);
829 assert_eq!(
830 result.symbols.len(),
831 0,
832 "lowercase non-typed assignment should not be extracted"
833 );
834 }
835
836 #[test]
837 fn test_class_level_assignment_skipped() {
838 let result = parse("class Foo:\n MAX = 100\n");
839 assert!(!result.has_errors);
840 let names: Vec<&str> = result.symbols.iter().map(|s| s.name.as_str()).collect();
842 assert!(names.contains(&"Foo"), "class should be extracted");
843 assert!(
844 !names.contains(&"MAX"),
845 "class-level constant should NOT be extracted, got: {:?}",
846 names
847 );
848 }
849
850 #[test]
851 fn test_method_vs_function() {
852 let result =
853 parse("def top():\n pass\n\nclass C:\n def inner(self):\n pass\n");
854 assert!(!result.has_errors);
855 let top = result.symbols.iter().find(|s| s.name == "top").unwrap();
856 let inner = result.symbols.iter().find(|s| s.name == "inner").unwrap();
857 assert_eq!(top.kind, SymbolKind::Function);
858 assert_eq!(inner.kind, SymbolKind::Method);
859 }
860
861 #[test]
862 fn test_empty_file() {
863 let result = parse("");
864 assert!(!result.has_errors);
865 assert_eq!(result.symbols.len(), 0);
866 }
867
868 #[test]
869 fn test_malformed_file() {
870 let result = parse("def broken(\n {}\n");
871 assert!(result.has_errors);
872 }
873
874 #[test]
877 fn test_import_simple() {
878 let result = parse("import os\n");
879 assert_eq!(result.imports.len(), 1);
880 assert_eq!(result.imports[0].path, "os");
881 assert_eq!(result.imports[0].imported_symbols, vec!["os"]);
882 }
883
884 #[test]
885 fn test_import_dotted() {
886 let result = parse("import os.path\n");
887 assert_eq!(result.imports.len(), 1);
888 assert_eq!(result.imports[0].path, "os.path");
889 assert_eq!(result.imports[0].imported_symbols, vec!["path"]);
890 }
891
892 #[test]
893 fn test_import_aliased() {
894 let result = parse("import numpy as np\n");
895 assert_eq!(result.imports.len(), 1);
896 assert_eq!(result.imports[0].path, "numpy");
897 assert_eq!(result.imports[0].imported_symbols, vec!["np"]);
898 }
899
900 #[test]
901 fn test_from_import() {
902 let result = parse("from os import path\n");
903 assert_eq!(result.imports.len(), 1);
904 assert_eq!(result.imports[0].path, "os");
905 assert!(
906 result.imports[0]
907 .imported_symbols
908 .contains(&"path".to_string()),
909 "imports: {:?}",
910 result.imports[0]
911 );
912 }
913
914 #[test]
915 fn test_from_import_multiple() {
916 let result = parse("from os.path import join, exists\n");
917 assert_eq!(result.imports.len(), 1);
918 assert_eq!(result.imports[0].path, "os.path");
919 assert!(result.imports[0]
920 .imported_symbols
921 .contains(&"join".to_string()));
922 assert!(result.imports[0]
923 .imported_symbols
924 .contains(&"exists".to_string()));
925 }
926
927 #[test]
928 fn test_from_import_relative() {
929 let result = parse("from . import utils\n");
930 assert_eq!(result.imports.len(), 1);
931 assert!(
933 result.imports[0]
934 .imported_symbols
935 .contains(&"utils".to_string()),
936 "imports: {:?}",
937 result.imports[0]
938 );
939 }
940
941 #[test]
942 fn test_from_import_wildcard() {
943 let result = parse("from foo import *\n");
944 assert_eq!(result.imports.len(), 1);
945 assert_eq!(result.imports[0].path, "foo");
946 assert!(
947 result.imports[0].imported_symbols.is_empty(),
948 "wildcard should have empty imported_symbols"
949 );
950 }
951
952 #[test]
955 fn test_function_call() {
956 let result = parse("def main():\n helper()\n\ndef helper():\n pass\n");
957 let call = result
958 .calls
959 .iter()
960 .find(|c| c.callee_name == "helper")
961 .expect("should find call to helper");
962 assert_eq!(call.caller_name, "main");
963 }
964
965 #[test]
966 fn test_method_call() {
967 let result = parse("def main():\n obj.process()\n");
968 let call = result
969 .calls
970 .iter()
971 .find(|c| c.callee_name == "process")
972 .expect("should find call to process");
973 assert_eq!(call.caller_name, "main");
974 }
975
976 #[test]
977 fn test_chained_method_call() {
978 let result = parse("def main():\n obj.foo().bar()\n");
979 let names: Vec<&str> = result
980 .calls
981 .iter()
982 .map(|c| c.callee_name.as_str())
983 .collect();
984 assert!(names.contains(&"foo"), "calls: {:?}", names);
985 assert!(names.contains(&"bar"), "calls: {:?}", names);
986 }
987
988 #[test]
989 fn test_builtin_call_filtered() {
990 let result = parse("def main():\n print(len(x))\n");
991 assert!(
992 result.calls.is_empty(),
993 "builtin calls should be filtered, got: {:?}",
994 result.calls
995 );
996 }
997
998 #[test]
1001 fn test_type_ref_parameter() {
1002 let result = parse("def foo(x: Config):\n pass\n");
1003 let refs: Vec<&str> = result
1004 .type_refs
1005 .iter()
1006 .map(|r| r.to_type.as_str())
1007 .collect();
1008 assert!(
1009 refs.contains(&"Config"),
1010 "should find Config type ref, got: {:?}",
1011 refs
1012 );
1013 }
1014
1015 #[test]
1016 fn test_type_ref_return() {
1017 let result = parse("def foo() -> Result:\n pass\n");
1018 let refs: Vec<&str> = result
1019 .type_refs
1020 .iter()
1021 .map(|r| r.to_type.as_str())
1022 .collect();
1023 assert!(
1024 refs.contains(&"Result"),
1025 "should find Result type ref, got: {:?}",
1026 refs
1027 );
1028 }
1029
1030 #[test]
1031 fn test_type_ref_assignment() {
1032 let result = parse("def foo():\n x: Config = get()\n");
1033 let refs: Vec<&str> = result
1034 .type_refs
1035 .iter()
1036 .map(|r| r.to_type.as_str())
1037 .collect();
1038 assert!(
1039 refs.contains(&"Config"),
1040 "should find Config type ref, got: {:?}",
1041 refs
1042 );
1043 }
1044
1045 #[test]
1046 fn test_builtin_type_filtered() {
1047 let result = parse("def foo(x: int, y: str) -> bool:\n pass\n");
1048 assert!(
1049 result.type_refs.is_empty(),
1050 "builtin types should be filtered, got: {:?}",
1051 result.type_refs
1052 );
1053 }
1054
1055 fn setup_python_project() -> tempfile::TempDir {
1058 let tmp = tempfile::tempdir().unwrap();
1059 let foo = tmp.path().join("foo");
1069 let utils = foo.join("utils");
1070 let src = tmp.path().join("src");
1071 std::fs::create_dir_all(&utils).unwrap();
1072 std::fs::create_dir_all(&src).unwrap();
1073 std::fs::write(foo.join("__init__.py"), "").unwrap();
1074 std::fs::write(foo.join("bar.py"), "def greet(): pass").unwrap();
1075 std::fs::write(utils.join("__init__.py"), "").unwrap();
1076 std::fs::write(utils.join("helpers.py"), "def help(): pass").unwrap();
1077 std::fs::write(src.join("app.py"), "def run(): pass").unwrap();
1078 std::fs::write(tmp.path().join("main.py"), "import foo").unwrap();
1079 tmp
1080 }
1081
1082 #[test]
1085 fn test_resolve_absolute_simple() {
1086 let tmp = setup_python_project();
1087 let result = resolve_python_import("foo.bar", "main.py", tmp.path());
1088 assert!(result.is_some(), "should resolve foo.bar");
1089 assert!(result.unwrap().ends_with("bar.py"));
1090 }
1091
1092 #[test]
1093 fn test_resolve_absolute_dotted_package() {
1094 let tmp = setup_python_project();
1095 let result = resolve_python_import("foo.utils", "main.py", tmp.path());
1096 assert!(result.is_some(), "should resolve foo.utils to __init__.py");
1097 let path = result.unwrap();
1098 assert!(path.ends_with("__init__.py"));
1099 assert!(path.to_string_lossy().contains("utils"));
1100 }
1101
1102 #[test]
1103 fn test_resolve_absolute_from_src() {
1104 let tmp = setup_python_project();
1105 let result = resolve_python_import("app", "main.py", tmp.path());
1106 assert!(result.is_some(), "should resolve app from src/ fallback");
1107 assert!(result.unwrap().ends_with("app.py"));
1108 }
1109
1110 #[test]
1111 fn test_resolve_relative_single_dot() {
1112 let tmp = setup_python_project();
1113 let result = resolve_python_import(".bar", "foo/__init__.py", tmp.path());
1114 assert!(result.is_some(), "should resolve .bar from foo/");
1115 assert!(result.unwrap().ends_with("bar.py"));
1116 }
1117
1118 #[test]
1119 fn test_resolve_relative_double_dot() {
1120 let tmp = setup_python_project();
1121 let result = resolve_python_import("..bar", "foo/utils/helpers.py", tmp.path());
1122 assert!(result.is_some(), "should resolve ..bar from foo/utils/");
1123 assert!(result.unwrap().ends_with("bar.py"));
1124 }
1125
1126 #[test]
1127 fn test_resolve_package_init() {
1128 let tmp = setup_python_project();
1129 let result = resolve_python_import("foo", "main.py", tmp.path());
1130 assert!(result.is_some(), "should resolve foo to __init__.py");
1131 assert!(result.unwrap().ends_with("__init__.py"));
1132 }
1133
1134 #[test]
1135 fn test_resolve_external_returns_none() {
1136 let tmp = setup_python_project();
1137 let result = resolve_python_import("numpy", "main.py", tmp.path());
1138 assert!(result.is_none(), "external package should return None");
1139 }
1140
1141 #[test]
1142 fn test_resolve_empty_path() {
1143 let tmp = setup_python_project();
1144 let result = resolve_python_import("", "main.py", tmp.path());
1145 assert!(result.is_none());
1146 }
1147
1148 #[test]
1149 fn test_resolve_relative_dot_init() {
1150 let tmp = setup_python_project();
1151 std::fs::write(tmp.path().join("foo/__init__.py"), "from . import utils").unwrap();
1153 let result = resolve_python_import(".utils", "foo/__init__.py", tmp.path());
1154 assert!(result.is_some(), "should resolve .utils from foo/");
1155 }
1156
1157 #[test]
1158 fn test_resolve_deep_dotted() {
1159 let tmp = setup_python_project();
1160 let result = resolve_python_import("foo.utils.helpers", "main.py", tmp.path());
1161 assert!(result.is_some(), "should resolve foo.utils.helpers");
1162 assert!(result.unwrap().ends_with("helpers.py"));
1163 }
1164
1165 #[test]
1168 fn test_decorated_class() {
1169 let result = parse(
1170 "@dataclass\nclass Config:\n name: str\n def validate(self):\n pass\n",
1171 );
1172 assert!(!result.has_errors);
1173 let config = result.symbols.iter().find(|s| s.name == "Config");
1174 assert!(config.is_some(), "decorated class should be extracted");
1175 assert_eq!(config.unwrap().kind, SymbolKind::Class);
1176 let sig = config.unwrap().signature.as_ref().unwrap();
1177 assert!(
1178 sig.contains("@dataclass"),
1179 "sig should contain decorator: {}",
1180 sig
1181 );
1182 assert!(
1183 sig.contains("class Config"),
1184 "sig should contain class def: {}",
1185 sig
1186 );
1187 let validate = result.symbols.iter().find(|s| s.name == "validate");
1189 assert!(
1190 validate.is_some(),
1191 "method in decorated class should be extracted"
1192 );
1193 assert_eq!(validate.unwrap().kind, SymbolKind::Method);
1194 }
1195
1196 #[test]
1197 fn test_decorated_class_with_methods_parent_index() {
1198 let result = parse("@decorator\nclass Svc:\n def run(self):\n pass\n def stop(self):\n pass\n");
1199 assert!(!result.has_errors);
1200 let svc = result.symbols.iter().find(|s| s.name == "Svc").unwrap();
1201 assert_eq!(svc.kind, SymbolKind::Class);
1202 let svc_idx = result.symbols.iter().position(|s| s.name == "Svc").unwrap();
1203 for method in result
1204 .symbols
1205 .iter()
1206 .filter(|s| s.kind == SymbolKind::Method)
1207 {
1208 assert_eq!(
1209 method.parent_index,
1210 Some(svc_idx),
1211 "method '{}' should have Svc as parent",
1212 method.name
1213 );
1214 }
1215 }
1216
1217 #[test]
1218 fn test_resolve_relative_dot_only_to_init() {
1219 let tmp = setup_python_project();
1221 let result = resolve_python_import(".", "foo/bar.py", tmp.path());
1222 assert!(
1223 result.is_some(),
1224 "bare dot should resolve to __init__.py in foo/"
1225 );
1226 assert!(result.unwrap().ends_with("__init__.py"));
1227 }
1228
1229 #[test]
1230 fn test_resolve_relative_dot_only_no_init() {
1231 let tmp = tempfile::tempdir().unwrap();
1233 let pkg = tmp.path().join("pkg");
1234 std::fs::create_dir_all(&pkg).unwrap();
1235 std::fs::write(pkg.join("mod.py"), "").unwrap();
1236 let result = resolve_python_import(".", "pkg/mod.py", tmp.path());
1238 assert!(
1239 result.is_none(),
1240 "bare dot without __init__.py should return None"
1241 );
1242 }
1243
1244 #[test]
1245 fn test_resolve_module_path_empty_parts() {
1246 let tmp = tempfile::tempdir().unwrap();
1248 let pkg = tmp.path().join("pkg");
1249 std::fs::create_dir_all(&pkg).unwrap();
1250 std::fs::write(pkg.join("__init__.py"), "").unwrap();
1251 let result = resolve_module_path(&[], &pkg);
1252 assert!(
1253 result.is_some(),
1254 "empty parts with __init__.py should resolve"
1255 );
1256 assert!(result.unwrap().ends_with("__init__.py"));
1257 }
1258
1259 #[test]
1260 fn test_resolve_module_path_empty_parts_no_init() {
1261 let tmp = tempfile::tempdir().unwrap();
1262 let pkg = tmp.path().join("pkg");
1263 std::fs::create_dir_all(&pkg).unwrap();
1264 let result = resolve_module_path(&[], &pkg);
1265 assert!(
1266 result.is_none(),
1267 "empty parts without __init__.py should return None"
1268 );
1269 }
1270
1271 #[test]
1272 fn test_resolve_intermediate_not_dir() {
1273 let tmp = setup_python_project();
1275 let result = resolve_python_import("foo.bar.nonexistent", "main.py", tmp.path());
1276 assert!(result.is_none(), "intermediate non-directory should fail");
1277 }
1278
1279 #[test]
1280 fn test_import_aliased_from() {
1281 let result = parse("from os import path as p\n");
1283 assert_eq!(result.imports.len(), 1);
1284 assert_eq!(result.imports[0].path, "os");
1285 assert!(
1286 result.imports[0]
1287 .imported_symbols
1288 .contains(&"p".to_string()),
1289 "aliased import should use alias name: {:?}",
1290 result.imports[0].imported_symbols
1291 );
1292 }
1293
1294 #[test]
1295 fn test_import_from_aliased_without_alias() {
1296 let result = parse("from os.path import join, exists\n");
1298 assert_eq!(result.imports.len(), 1);
1299 assert!(result.imports[0]
1300 .imported_symbols
1301 .contains(&"join".to_string()));
1302 assert!(result.imports[0]
1303 .imported_symbols
1304 .contains(&"exists".to_string()));
1305 }
1306
1307 #[test]
1308 fn test_multiple_plain_imports() {
1309 let result = parse("import os\nimport sys\nimport json\n");
1310 assert_eq!(result.imports.len(), 3);
1311 let paths: Vec<&str> = result.imports.iter().map(|i| i.path.as_str()).collect();
1312 assert!(paths.contains(&"os"));
1313 assert!(paths.contains(&"sys"));
1314 assert!(paths.contains(&"json"));
1315 }
1316
1317 #[test]
1318 fn test_triple_dot_relative_import() {
1319 let tmp = tempfile::tempdir().unwrap();
1320 let deep = tmp.path().join("a").join("b").join("c");
1321 std::fs::create_dir_all(&deep).unwrap();
1322 std::fs::write(deep.join("mod.py"), "").unwrap();
1323 std::fs::write(tmp.path().join("a").join("target.py"), "").unwrap();
1324 let result = resolve_python_import("...target", "a/b/c/mod.py", tmp.path());
1325 assert!(result.is_some(), "triple dot should resolve two levels up");
1326 assert!(result.unwrap().ends_with("target.py"));
1327 }
1328
1329 #[test]
1330 fn test_resolve_nonexistent_module() {
1331 let tmp = setup_python_project();
1332 let result = resolve_python_import("foo.nonexistent", "main.py", tmp.path());
1333 assert!(result.is_none(), "nonexistent module should return None");
1334 }
1335
1336 #[test]
1337 fn test_call_with_attribute_access() {
1338 let result = parse("def main():\n config.get_value()\n");
1340 let calls: Vec<&str> = result
1341 .calls
1342 .iter()
1343 .map(|c| c.callee_name.as_str())
1344 .collect();
1345 assert!(
1346 calls.contains(&"get_value"),
1347 "should extract attribute call: {:?}",
1348 calls
1349 );
1350 }
1351
1352 #[test]
1353 fn test_type_ref_builtin_type_not_extracted() {
1354 let result = parse("def foo(x: int, y: str, z: list) -> bool:\n pass\n");
1356 assert!(
1357 result.type_refs.is_empty(),
1358 "builtin types should be filtered: {:?}",
1359 result.type_refs
1360 );
1361 }
1362
1363 #[test]
1364 fn test_resolve_from_lib_directory() {
1365 let tmp = tempfile::tempdir().unwrap();
1367 let lib = tmp.path().join("lib");
1368 std::fs::create_dir_all(&lib).unwrap();
1369 std::fs::write(lib.join("mymod.py"), "").unwrap();
1370 let result = resolve_python_import("mymod", "main.py", tmp.path());
1371 assert!(result.is_some(), "should resolve from lib/ fallback");
1372 assert!(result.unwrap().ends_with("mymod.py"));
1373 }
1374
1375 #[test]
1376 fn test_expression_statement_not_assignment() {
1377 let result = parse("print('hello')\n");
1379 assert!(!result.has_errors);
1380 assert_eq!(
1381 result.symbols.len(),
1382 0,
1383 "bare expression should not produce symbols"
1384 );
1385 }
1386
1387 #[test]
1388 fn test_assignment_left_not_identifier() {
1389 let result = parse("x, y = 1, 2\n");
1391 assert!(!result.has_errors);
1392 }
1394
1395 #[test]
1396 fn test_is_python_builtin_positive() {
1397 assert!(is_python_builtin("print"));
1398 assert!(is_python_builtin("len"));
1399 assert!(is_python_builtin("isinstance"));
1400 assert!(is_python_builtin("ValueError"));
1401 assert!(is_python_builtin("super"));
1402 assert!(is_python_builtin("object"));
1403 }
1404
1405 #[test]
1406 fn test_is_python_builtin_negative() {
1407 assert!(!is_python_builtin("my_function"));
1408 assert!(!is_python_builtin("CustomError"));
1409 assert!(!is_python_builtin(""));
1410 }
1411
1412 #[test]
1413 fn test_is_python_builtin_type_positive() {
1414 assert!(is_python_builtin_type("int"));
1415 assert!(is_python_builtin_type("str"));
1416 assert!(is_python_builtin_type("Optional"));
1417 assert!(is_python_builtin_type("List"));
1418 assert!(is_python_builtin_type("TypeVar"));
1419 assert!(is_python_builtin_type("Self"));
1420 assert!(is_python_builtin_type("Never"));
1421 assert!(is_python_builtin_type("NoReturn"));
1422 }
1423
1424 #[test]
1425 fn test_is_python_builtin_type_negative() {
1426 assert!(!is_python_builtin_type("MyClass"));
1427 assert!(!is_python_builtin_type("AppConfig"));
1428 assert!(!is_python_builtin_type(""));
1429 }
1430
1431 #[test]
1432 fn test_type_ref_inside_class_method() {
1433 let result =
1434 parse("class Svc:\n def process(self, cfg: AppConfig) -> Response:\n pass\n");
1435 assert!(!result.has_errors);
1436 let ref_types: Vec<&str> = result
1437 .type_refs
1438 .iter()
1439 .map(|r| r.to_type.as_str())
1440 .collect();
1441 assert!(
1442 ref_types.contains(&"AppConfig"),
1443 "should find AppConfig type ref: {:?}",
1444 ref_types
1445 );
1446 assert!(
1447 ref_types.contains(&"Response"),
1448 "should find Response type ref: {:?}",
1449 ref_types
1450 );
1451 }
1452
1453 #[test]
1454 fn test_module_constant_underscore_allcaps() {
1455 let result = parse("MAX_RETRY_COUNT = 3\n");
1456 assert!(!result.has_errors);
1457 assert_eq!(result.symbols.len(), 1);
1458 assert_eq!(result.symbols[0].name, "MAX_RETRY_COUNT");
1459 assert_eq!(result.symbols[0].kind, SymbolKind::Variable);
1460 }
1461
1462 #[test]
1463 fn test_resolve_absolute_nonexistent_package() {
1464 let tmp = setup_python_project();
1465 let result = resolve_python_import("totally_missing.submod", "main.py", tmp.path());
1466 assert!(result.is_none());
1467 }
1468
1469 #[test]
1470 fn test_resolve_relative_too_many_dots() {
1471 let tmp = tempfile::tempdir().unwrap();
1472 let pkg = tmp.path().join("a");
1473 std::fs::create_dir_all(&pkg).unwrap();
1474 std::fs::write(pkg.join("mod.py"), "").unwrap();
1475 let result = resolve_python_import("....target", "a/mod.py", tmp.path());
1476 assert!(result.is_none(), "too many dots should return None");
1477 }
1478
1479 #[test]
1480 fn test_import_plain_multiple_dotted_names() {
1481 let result = parse("import os, sys, json\n");
1483 assert!(!result.has_errors);
1484 let paths: Vec<&str> = result.imports.iter().map(|i| i.path.as_str()).collect();
1485 assert!(paths.contains(&"os"), "imports: {:?}", paths);
1486 assert!(paths.contains(&"sys"), "imports: {:?}", paths);
1487 assert!(paths.contains(&"json"), "imports: {:?}", paths);
1488 }
1489
1490 #[test]
1491 fn test_from_import_with_alias() {
1492 let result = parse("from os import path as p\n");
1493 assert_eq!(result.imports.len(), 1);
1494 assert_eq!(result.imports[0].path, "os");
1495 assert!(
1496 result.imports[0]
1497 .imported_symbols
1498 .contains(&"p".to_string()),
1499 "aliased import should use alias name: {:?}",
1500 result.imports[0].imported_symbols
1501 );
1502 }
1503
1504 #[test]
1505 fn test_call_identifier_branch() {
1506 let result = parse("def main():\n do_work()\n");
1508 assert!(!result.has_errors);
1509 let calls: Vec<&str> = result
1510 .calls
1511 .iter()
1512 .map(|c| c.callee_name.as_str())
1513 .collect();
1514 assert!(
1515 calls.contains(&"do_work"),
1516 "should capture identifier call: {:?}",
1517 calls
1518 );
1519 }
1520
1521 #[test]
1522 fn test_decorated_definition_no_definition_node() {
1523 let result = parse("@property\ndef x(self):\n return 1\n");
1527 assert!(!result.has_errors);
1528 assert_eq!(result.symbols.len(), 1);
1529 assert_eq!(result.symbols[0].name, "x");
1530 }
1531
1532 #[test]
1535 fn test_class_with_inheritance() {
1536 let result = parse("class Child(Parent):\n def method(self):\n pass\n");
1537 assert!(!result.has_errors);
1538 let child = result.symbols.iter().find(|s| s.name == "Child").unwrap();
1539 assert_eq!(child.kind, SymbolKind::Class);
1540 let sig = child.signature.as_ref().unwrap();
1541 assert!(
1542 sig.contains("Parent"),
1543 "sig should contain base class: {}",
1544 sig
1545 );
1546 }
1547
1548 #[test]
1549 fn test_class_multiple_inheritance() {
1550 let result = parse("class Multi(Base1, Base2, Mixin):\n pass\n");
1551 assert!(!result.has_errors);
1552 assert_eq!(result.symbols.len(), 1);
1553 assert_eq!(result.symbols[0].name, "Multi");
1554 let sig = result.symbols[0].signature.as_ref().unwrap();
1555 assert!(sig.contains("Base1"), "sig = {}", sig);
1556 assert!(sig.contains("Mixin"), "sig = {}", sig);
1557 }
1558
1559 #[test]
1560 fn test_property_decorated_method() {
1561 let result =
1562 parse("class Foo:\n @property\n def name(self):\n return self._name\n");
1563 assert!(!result.has_errors);
1564 let name_method = result.symbols.iter().find(|s| s.name == "name").unwrap();
1565 assert_eq!(name_method.kind, SymbolKind::Method);
1566 let sig = name_method.signature.as_ref().unwrap();
1567 assert!(sig.contains("@property"), "sig = {}", sig);
1568 }
1569
1570 #[test]
1571 fn test_classmethod_decorated() {
1572 let result = parse("class Foo:\n @classmethod\n def create(cls):\n pass\n");
1573 assert!(!result.has_errors);
1574 let create = result.symbols.iter().find(|s| s.name == "create").unwrap();
1575 assert_eq!(create.kind, SymbolKind::Method);
1576 let sig = create.signature.as_ref().unwrap();
1577 assert!(sig.contains("@classmethod"), "sig = {}", sig);
1578 }
1579
1580 #[test]
1581 fn test_function_with_complex_signature() {
1582 let result = parse(
1583 "def process(items: list[str], *, timeout: int = 30) -> dict[str, Any]:\n pass\n",
1584 );
1585 assert!(!result.has_errors);
1586 assert_eq!(result.symbols.len(), 1);
1587 assert_eq!(result.symbols[0].name, "process");
1588 let sig = result.symbols[0].signature.as_ref().unwrap();
1589 assert!(sig.contains("def process"), "sig = {}", sig);
1590 }
1591
1592 #[test]
1593 fn test_call_inside_method() {
1594 let result =
1595 parse("class Svc:\n def run(self):\n self.setup()\n process()\n");
1596 let calls: Vec<&str> = result
1597 .calls
1598 .iter()
1599 .map(|c| c.callee_name.as_str())
1600 .collect();
1601 assert!(
1602 calls.contains(&"setup"),
1603 "should find self.setup(): {:?}",
1604 calls
1605 );
1606 assert!(
1607 calls.contains(&"process"),
1608 "should find process(): {:?}",
1609 calls
1610 );
1611 }
1612
1613 #[test]
1614 fn test_type_ref_generic_annotation() {
1615 let result = parse("def foo(items: List[Config]) -> Optional[Response]:\n pass\n");
1616 let refs: Vec<&str> = result
1617 .type_refs
1618 .iter()
1619 .map(|r| r.to_type.as_str())
1620 .collect();
1621 assert!(refs.contains(&"Config"), "type_refs: {:?}", refs);
1622 assert!(refs.contains(&"Response"), "type_refs: {:?}", refs);
1623 }
1624
1625 #[test]
1626 fn test_multiple_functions_and_classes() {
1627 let result = parse(
1628 "def a():\n pass\ndef b():\n pass\nclass C:\n def d(self):\n pass\n",
1629 );
1630 assert!(!result.has_errors);
1631 let names: Vec<&str> = result.symbols.iter().map(|s| s.name.as_str()).collect();
1632 assert_eq!(names, vec!["a", "b", "C", "d"]);
1633 }
1634
1635 #[test]
1636 fn test_module_constant_with_numbers_not_extracted() {
1637 let result = parse("MAX_SIZE_3 = 1024\n");
1639 assert!(!result.has_errors);
1640 assert_eq!(
1641 result.symbols.len(),
1642 0,
1643 "allcaps with digits should NOT be extracted (digits fail the check)"
1644 );
1645 }
1646
1647 #[test]
1648 fn test_module_constant_pure_allcaps() {
1649 let result = parse("MAX_SIZE = 1024\n");
1650 assert!(!result.has_errors);
1651 assert_eq!(result.symbols.len(), 1);
1652 assert_eq!(result.symbols[0].name, "MAX_SIZE");
1653 }
1654
1655 #[test]
1656 fn test_from_import_relative_double_dot() {
1657 let result = parse("from .. import utils\n");
1658 assert!(!result.has_errors);
1659 assert_eq!(result.imports.len(), 1);
1660 assert!(result.imports[0]
1661 .imported_symbols
1662 .contains(&"utils".to_string()));
1663 }
1664
1665 #[test]
1666 fn test_language_parser_trait_name() {
1667 let parser = PythonParser;
1668 assert_eq!(parser.language_name(), "python");
1669 }
1670
1671 #[test]
1672 fn test_language_parser_parse_via_trait() {
1673 let parser = PythonParser;
1674 let result = parser.parse("def foo():\n pass\n", &PathBuf::from("test.py"));
1675 assert!(!result.has_errors);
1676 assert_eq!(result.symbols.len(), 1);
1677 assert_eq!(result.symbols[0].name, "foo");
1678 }
1679
1680 #[test]
1681 fn test_resolve_from_init_subpackage() {
1682 let tmp = setup_python_project();
1683 let result = resolve_python_import("foo.utils.helpers", "foo/__init__.py", tmp.path());
1684 assert!(result.is_some(), "should resolve deep submodule from init");
1685 assert!(result.unwrap().ends_with("helpers.py"));
1686 }
1687
1688 #[test]
1689 fn test_call_nested_in_class_method() {
1690 let result =
1691 parse("class A:\n def m(self):\n x = helper()\n y = other()\n");
1692 let callers: Vec<&str> = result
1693 .calls
1694 .iter()
1695 .map(|c| c.caller_name.as_str())
1696 .collect();
1697 for caller in &callers {
1699 assert_eq!(*caller, "m");
1700 }
1701 }
1702}