1use crate::models::{Language, SearchResult, Span, SymbolKind};
15use anyhow::{Context, Result};
16use streaming_iterator::StreamingIterator;
17use tree_sitter::{Parser, Query, QueryCursor};
18
19pub fn parse(path: &str, source: &str) -> Result<Vec<SearchResult>> {
21 let mut parser = Parser::new();
22 let language = tree_sitter_python::LANGUAGE;
23
24 parser
25 .set_language(&language.into())
26 .context("Failed to set Python language")?;
27
28 let tree = parser
29 .parse(source, None)
30 .context("Failed to parse Python source")?;
31
32 let root_node = tree.root_node();
33
34 let mut symbols = Vec::new();
35
36 symbols.extend(extract_functions(source, &root_node, &language.into())?);
38 symbols.extend(extract_classes(source, &root_node, &language.into())?);
39 symbols.extend(extract_methods(source, &root_node, &language.into())?);
40 symbols.extend(extract_constants(source, &root_node, &language.into())?);
41 symbols.extend(extract_global_variables(
42 source,
43 &root_node,
44 &language.into(),
45 )?);
46 symbols.extend(extract_local_variables(
47 source,
48 &root_node,
49 &language.into(),
50 )?);
51 symbols.extend(extract_lambdas(source, &root_node, &language.into())?);
52
53 for symbol in &mut symbols {
55 symbol.path = path.to_string();
56 symbol.lang = Language::Python;
57 }
58
59 Ok(symbols)
60}
61
62fn extract_functions(
64 source: &str,
65 root: &tree_sitter::Node,
66 language: &tree_sitter::Language,
67) -> Result<Vec<SearchResult>> {
68 let query_str = r#"
69 (function_definition
70 name: (identifier) @name) @function
71 "#;
72
73 let query = Query::new(language, query_str).context("Failed to create function query")?;
74
75 extract_symbols(source, root, &query, SymbolKind::Function, None)
76}
77
78fn extract_classes(
80 source: &str,
81 root: &tree_sitter::Node,
82 language: &tree_sitter::Language,
83) -> Result<Vec<SearchResult>> {
84 let query_str = r#"
85 (class_definition
86 name: (identifier) @name) @class
87 "#;
88
89 let query = Query::new(language, query_str).context("Failed to create class query")?;
90
91 extract_symbols(source, root, &query, SymbolKind::Class, None)
92}
93
94fn extract_methods(
96 source: &str,
97 root: &tree_sitter::Node,
98 language: &tree_sitter::Language,
99) -> Result<Vec<SearchResult>> {
100 let query_str = r#"
101 (class_definition
102 name: (identifier) @class_name
103 body: (block
104 (function_definition
105 name: (identifier) @method_name))) @class
106
107 (class_definition
108 name: (identifier) @class_name
109 body: (block
110 (decorated_definition
111 (function_definition
112 name: (identifier) @method_name)))) @class
113 "#;
114
115 let query = Query::new(language, query_str).context("Failed to create method query")?;
116
117 let mut cursor = QueryCursor::new();
118 let mut matches = cursor.matches(&query, *root, source.as_bytes());
119
120 let mut symbols = Vec::new();
121
122 while let Some(match_) = matches.next() {
123 let mut class_name = None;
124 let mut method_name = None;
125 let mut method_node = None;
126
127 for capture in match_.captures {
128 let capture_name: &str = query.capture_names()[capture.index as usize];
129 match capture_name {
130 "class_name" => {
131 class_name = Some(
132 capture
133 .node
134 .utf8_text(source.as_bytes())
135 .unwrap_or("")
136 .to_string(),
137 );
138 }
139 "method_name" => {
140 method_name = Some(
141 capture
142 .node
143 .utf8_text(source.as_bytes())
144 .unwrap_or("")
145 .to_string(),
146 );
147 let mut current = capture.node;
149 while let Some(parent) = current.parent() {
150 if parent.kind() == "function_definition" {
151 method_node = Some(parent);
152 break;
153 }
154 current = parent;
155 }
156 }
157 _ => {}
158 }
159 }
160
161 if let (Some(class_name), Some(method_name), Some(node)) =
162 (class_name, method_name, method_node)
163 {
164 let scope = format!("class {}", class_name);
165 let span = node_to_span(&node);
166 let preview = extract_preview(source, &span);
167
168 symbols.push(SearchResult::new(
169 String::new(),
170 Language::Python,
171 SymbolKind::Method,
172 Some(method_name),
173 span,
174 Some(scope),
175 preview,
176 ));
177 }
178 }
179
180 Ok(symbols)
181}
182
183fn extract_constants(
185 source: &str,
186 root: &tree_sitter::Node,
187 language: &tree_sitter::Language,
188) -> Result<Vec<SearchResult>> {
189 let query_str = r#"
190 (module
191 (expression_statement
192 (assignment
193 left: (identifier) @name))) @const
194 "#;
195
196 let query = Query::new(language, query_str).context("Failed to create constant query")?;
197
198 let mut cursor = QueryCursor::new();
199 let mut matches = cursor.matches(&query, *root, source.as_bytes());
200
201 let mut symbols = Vec::new();
202
203 while let Some(match_) = matches.next() {
204 let mut name = None;
205 let mut const_node = None;
206
207 for capture in match_.captures {
208 let capture_name: &str = query.capture_names()[capture.index as usize];
209 if capture_name == "name" {
210 let name_text = capture.node.utf8_text(source.as_bytes()).unwrap_or("");
211 if name_text
213 .chars()
214 .all(|c| c.is_uppercase() || c == '_' || c.is_numeric())
215 {
216 name = Some(name_text.to_string());
217 let mut current = capture.node;
219 while let Some(parent) = current.parent() {
220 if parent.kind() == "assignment" {
221 const_node = Some(parent);
222 break;
223 }
224 current = parent;
225 }
226 }
227 }
228 }
229
230 if let (Some(name), Some(node)) = (name, const_node) {
231 let span = node_to_span(&node);
232 let preview = extract_preview(source, &span);
233
234 symbols.push(SearchResult::new(
235 String::new(),
236 Language::Python,
237 SymbolKind::Constant,
238 Some(name),
239 span,
240 None,
241 preview,
242 ));
243 }
244 }
245
246 Ok(symbols)
247}
248
249fn extract_global_variables(
251 source: &str,
252 root: &tree_sitter::Node,
253 language: &tree_sitter::Language,
254) -> Result<Vec<SearchResult>> {
255 let query_str = r#"
256 (module
257 (expression_statement
258 (assignment
259 left: (identifier) @name))) @var
260 "#;
261
262 let query =
263 Query::new(language, query_str).context("Failed to create global variable query")?;
264
265 let mut cursor = QueryCursor::new();
266 let mut matches = cursor.matches(&query, *root, source.as_bytes());
267
268 let mut symbols = Vec::new();
269
270 while let Some(match_) = matches.next() {
271 let mut name = None;
272 let mut var_node = None;
273
274 for capture in match_.captures {
275 let capture_name: &str = query.capture_names()[capture.index as usize];
276 if capture_name == "name" {
277 let name_text = capture.node.utf8_text(source.as_bytes()).unwrap_or("");
278 if !name_text
280 .chars()
281 .all(|c| c.is_uppercase() || c == '_' || c.is_numeric())
282 {
283 name = Some(name_text.to_string());
284 let mut current = capture.node;
286 while let Some(parent) = current.parent() {
287 if parent.kind() == "assignment" {
288 var_node = Some(parent);
289 break;
290 }
291 current = parent;
292 }
293 }
294 }
295 }
296
297 if let (Some(name), Some(node)) = (name, var_node) {
298 let span = node_to_span(&node);
299 let preview = extract_preview(source, &span);
300
301 symbols.push(SearchResult::new(
302 String::new(),
303 Language::Python,
304 SymbolKind::Variable,
305 Some(name),
306 span,
307 None,
308 preview,
309 ));
310 }
311 }
312
313 Ok(symbols)
314}
315
316fn extract_local_variables(
318 source: &str,
319 root: &tree_sitter::Node,
320 language: &tree_sitter::Language,
321) -> Result<Vec<SearchResult>> {
322 let query_str = r#"
323 (assignment
324 left: (identifier) @name) @assignment
325 "#;
326
327 let query = Query::new(language, query_str).context("Failed to create local variable query")?;
328
329 let mut cursor = QueryCursor::new();
330 let mut matches = cursor.matches(&query, *root, source.as_bytes());
331
332 let mut symbols = Vec::new();
333
334 while let Some(match_) = matches.next() {
335 let mut name = None;
336 let mut assignment_node = None;
337
338 for capture in match_.captures {
339 let capture_name: &str = query.capture_names()[capture.index as usize];
340 match capture_name {
341 "name" => {
342 let name_text = capture.node.utf8_text(source.as_bytes()).unwrap_or("");
343 if !name_text
345 .chars()
346 .all(|c| c.is_uppercase() || c == '_' || c.is_numeric())
347 {
348 name = Some(name_text.to_string());
349 }
350 }
351 "assignment" => {
352 assignment_node = Some(capture.node);
353 }
354 _ => {}
355 }
356 }
357
358 if let (Some(name), Some(node)) = (name, assignment_node) {
360 let mut is_in_function = false;
361 let mut current = node;
362
363 while let Some(parent) = current.parent() {
364 if parent.kind() == "function_definition" {
365 is_in_function = true;
366 break;
367 }
368 if parent.kind() == "module" {
370 break;
371 }
372 current = parent;
373 }
374
375 if is_in_function {
376 let span = node_to_span(&node);
377 let preview = extract_preview(source, &span);
378
379 symbols.push(SearchResult::new(
380 String::new(),
381 Language::Python,
382 SymbolKind::Variable,
383 Some(name),
384 span,
385 None, preview,
387 ));
388 }
389 }
390 }
391
392 Ok(symbols)
393}
394
395fn extract_lambdas(
397 source: &str,
398 root: &tree_sitter::Node,
399 language: &tree_sitter::Language,
400) -> Result<Vec<SearchResult>> {
401 let query_str = r#"
402 (assignment
403 left: (identifier) @name
404 right: (lambda)) @lambda
405 "#;
406
407 let query = Query::new(language, query_str).context("Failed to create lambda query")?;
408
409 extract_symbols(source, root, &query, SymbolKind::Function, None)
410}
411
412fn extract_symbols(
414 source: &str,
415 root: &tree_sitter::Node,
416 query: &Query,
417 kind: SymbolKind,
418 scope: Option<String>,
419) -> Result<Vec<SearchResult>> {
420 let mut cursor = QueryCursor::new();
421 let mut matches = cursor.matches(query, *root, source.as_bytes());
422
423 let mut symbols = Vec::new();
424
425 while let Some(match_) = matches.next() {
426 let mut name = None;
428 let mut full_node = None;
429
430 for capture in match_.captures {
431 let capture_name: &str = query.capture_names()[capture.index as usize];
432 if capture_name == "name" {
433 name = Some(
434 capture
435 .node
436 .utf8_text(source.as_bytes())
437 .unwrap_or("")
438 .to_string(),
439 );
440 } else {
441 full_node = Some(capture.node);
443 }
444 }
445
446 if let (Some(name), Some(node)) = (name, full_node) {
447 let span = node_to_span(&node);
448 let preview = extract_preview(source, &span);
449
450 symbols.push(SearchResult::new(
451 String::new(),
452 Language::Python,
453 kind.clone(),
454 Some(name),
455 span,
456 scope.clone(),
457 preview,
458 ));
459 }
460 }
461
462 Ok(symbols)
463}
464
465fn node_to_span(node: &tree_sitter::Node) -> Span {
467 let start = node.start_position();
468 let end = node.end_position();
469
470 Span::new(
471 start.row + 1, start.column,
473 end.row + 1,
474 end.column,
475 )
476}
477
478fn extract_preview(source: &str, span: &Span) -> String {
480 let lines: Vec<&str> = source.lines().collect();
481
482 let start_idx = span.start_line - 1; let end_idx = (start_idx + 7).min(lines.len());
485
486 lines[start_idx..end_idx].join("\n")
487}
488
489use crate::models::ImportType;
494use crate::parsers::{DependencyExtractor, ImportInfo};
495
496pub struct PythonDependencyExtractor;
498
499impl DependencyExtractor for PythonDependencyExtractor {
500 fn extract_dependencies(source: &str) -> Result<Vec<ImportInfo>> {
501 let mut parser = Parser::new();
502 let language = tree_sitter_python::LANGUAGE;
503
504 parser
505 .set_language(&language.into())
506 .context("Failed to set Python language")?;
507
508 let tree = parser
509 .parse(source, None)
510 .context("Failed to parse Python source")?;
511
512 let root_node = tree.root_node();
513
514 let mut imports = Vec::new();
515
516 imports.extend(extract_import_statements(source, &root_node)?);
518
519 imports.extend(extract_from_imports(source, &root_node)?);
521
522 Ok(imports)
523 }
524}
525
526fn extract_import_statements(source: &str, root: &tree_sitter::Node) -> Result<Vec<ImportInfo>> {
528 let language = tree_sitter_python::LANGUAGE;
529
530 let query_str = r#"
531 (import_statement
532 name: (dotted_name) @import_path) @import
533 "#;
534
535 let query = Query::new(&language.into(), query_str)
536 .context("Failed to create import statement query")?;
537
538 let mut cursor = QueryCursor::new();
539 let mut matches = cursor.matches(&query, *root, source.as_bytes());
540
541 let mut imports = Vec::new();
542
543 while let Some(match_) = matches.next() {
544 let mut import_path = None;
545 let mut import_node = None;
546
547 for capture in match_.captures {
548 let capture_name: &str = query.capture_names()[capture.index as usize];
549 match capture_name {
550 "import_path" => {
551 import_path = Some(
552 capture
553 .node
554 .utf8_text(source.as_bytes())
555 .unwrap_or("")
556 .to_string(),
557 );
558 }
559 "import" => {
560 import_node = Some(capture.node);
561 }
562 _ => {}
563 }
564 }
565
566 if let (Some(path), Some(node)) = (import_path, import_node) {
567 let import_type = classify_python_import(&path);
568 let line_number = node.start_position().row + 1;
569
570 imports.push(ImportInfo {
571 imported_path: path,
572 import_type,
573 line_number,
574 imported_symbols: None,
575 });
576 }
577 }
578
579 Ok(imports)
580}
581
582fn extract_from_imports(source: &str, root: &tree_sitter::Node) -> Result<Vec<ImportInfo>> {
584 let language = tree_sitter_python::LANGUAGE;
585
586 let query_str = r#"
587 (import_from_statement
588 module_name: (dotted_name) @module_path) @import
589
590 (import_from_statement
591 module_name: (relative_import) @module_path) @import
592 "#;
593
594 let query =
595 Query::new(&language.into(), query_str).context("Failed to create from-import query")?;
596
597 let mut cursor = QueryCursor::new();
598 let mut matches = cursor.matches(&query, *root, source.as_bytes());
599
600 let mut imports = Vec::new();
601
602 while let Some(match_) = matches.next() {
603 let mut module_path = None;
604 let mut import_node = None;
605
606 for capture in match_.captures {
607 let capture_name: &str = query.capture_names()[capture.index as usize];
608 match capture_name {
609 "module_path" => {
610 module_path = Some(
611 capture
612 .node
613 .utf8_text(source.as_bytes())
614 .unwrap_or("")
615 .to_string(),
616 );
617 }
618 "import" => {
619 import_node = Some(capture.node);
620 }
621 _ => {}
622 }
623 }
624
625 if let (Some(path), Some(node)) = (module_path, import_node) {
626 let import_type = classify_python_import(&path);
627 let line_number = node.start_position().row + 1;
628
629 let imported_symbols = extract_imported_symbols(source, &node);
631
632 imports.push(ImportInfo {
633 imported_path: path,
634 import_type,
635 line_number,
636 imported_symbols,
637 });
638 }
639 }
640
641 Ok(imports)
642}
643
644fn extract_imported_symbols(source: &str, import_node: &tree_sitter::Node) -> Option<Vec<String>> {
646 let mut symbols = Vec::new();
647
648 let mut cursor = import_node.walk();
650 for child in import_node.children(&mut cursor) {
651 match child.kind() {
652 "aliased_import" | "dotted_name" => {
653 let mut child_cursor = child.walk();
655 for grandchild in child.children(&mut child_cursor) {
656 if (grandchild.kind() == "identifier" || grandchild.kind() == "dotted_name")
657 && let Ok(text) = grandchild.utf8_text(source.as_bytes())
658 {
659 symbols.push(text.to_string());
660 break; }
662 }
663 }
664 _ => {}
665 }
666 }
667
668 if symbols.is_empty() {
669 None
670 } else {
671 Some(symbols)
672 }
673}
674
675pub fn find_python_package_name(root: &std::path::Path) -> Option<String> {
678 if let Some(name) = find_pyproject_package(root) {
680 return Some(name);
681 }
682
683 if let Some(name) = find_setup_py_package(root) {
685 return Some(name);
686 }
687
688 if let Some(name) = find_setup_cfg_package(root) {
690 return Some(name);
691 }
692
693 None
694}
695
696fn find_pyproject_package(root: &std::path::Path) -> Option<String> {
698 let pyproject_path = root.join("pyproject.toml");
699 let content = std::fs::read_to_string(pyproject_path).ok()?;
700
701 let mut in_project_section = false;
704
705 for line in content.lines() {
706 let trimmed = line.trim();
707
708 if trimmed == "[project]" {
710 in_project_section = true;
711 continue;
712 }
713
714 if trimmed.starts_with('[') && trimmed != "[project]" {
716 in_project_section = false;
717 continue;
718 }
719
720 if in_project_section
722 && trimmed.starts_with("name")
723 && trimmed.contains('=')
724 && let Some(equals_pos) = trimmed.find('=')
725 {
726 let after_equals = trimmed[equals_pos + 1..].trim();
727
728 for quote in ['"', '\''] {
730 if let Some(start) = after_equals.find(quote)
731 && let Some(end) = after_equals[start + 1..].find(quote)
732 {
733 let name = &after_equals[start + 1..start + 1 + end];
734 return Some(name.to_lowercase());
736 }
737 }
738 }
739 }
740
741 None
742}
743
744fn find_setup_py_package(root: &std::path::Path) -> Option<String> {
746 let setup_path = root.join("setup.py");
747 let content = std::fs::read_to_string(setup_path).ok()?;
748
749 for line in content.lines() {
752 let trimmed = line.trim();
753
754 if trimmed.contains("name") && trimmed.contains('=') {
755 if let Some(name_pos) = trimmed.find("name") {
757 let after_name = &trimmed[name_pos + 4..]; if let Some(equals_pos) = after_name.find('=') {
760 let after_equals = after_name[equals_pos + 1..].trim();
761
762 for quote in ['"', '\''] {
764 if let Some(start) = after_equals.find(quote)
765 && let Some(end) = after_equals[start + 1..].find(quote)
766 {
767 let name = &after_equals[start + 1..start + 1 + end];
768 return Some(name.to_lowercase());
769 }
770 }
771 }
772 }
773 }
774 }
775
776 None
777}
778
779fn find_setup_cfg_package(root: &std::path::Path) -> Option<String> {
781 let setup_cfg_path = root.join("setup.cfg");
782 let content = std::fs::read_to_string(setup_cfg_path).ok()?;
783
784 let mut in_metadata_section = false;
786
787 for line in content.lines() {
788 let trimmed = line.trim();
789
790 if trimmed == "[metadata]" {
792 in_metadata_section = true;
793 continue;
794 }
795
796 if trimmed.starts_with('[') && trimmed != "[metadata]" {
798 in_metadata_section = false;
799 continue;
800 }
801
802 if in_metadata_section
804 && trimmed.starts_with("name")
805 && trimmed.contains('=')
806 && let Some(equals_pos) = trimmed.find('=')
807 {
808 let name = trimmed[equals_pos + 1..].trim();
809 return Some(name.to_lowercase());
810 }
811 }
812
813 None
814}
815
816pub fn reclassify_python_import(import_path: &str, package_prefix: Option<&str>) -> ImportType {
819 if let Some(prefix) = package_prefix {
821 let first_component = import_path.split('.').next().unwrap_or(import_path);
823
824 if first_component == prefix {
825 return ImportType::Internal;
826 }
827 }
828
829 if import_path.starts_with('.') {
831 return ImportType::Internal;
832 }
833
834 if is_python_stdlib(import_path) {
836 return ImportType::Stdlib;
837 }
838
839 ImportType::External
841}
842
843fn is_python_stdlib(path: &str) -> bool {
845 const STDLIB_MODULES: &[&str] = &[
846 "os",
847 "sys",
848 "io",
849 "re",
850 "json",
851 "csv",
852 "xml",
853 "html",
854 "http",
855 "urllib",
856 "collections",
857 "itertools",
858 "functools",
859 "operator",
860 "pathlib",
861 "glob",
862 "tempfile",
863 "shutil",
864 "pickle",
865 "shelve",
866 "sqlite3",
867 "zlib",
868 "gzip",
869 "time",
870 "datetime",
871 "calendar",
872 "logging",
873 "argparse",
874 "configparser",
875 "typing",
876 "dataclasses",
877 "enum",
878 "abc",
879 "contextlib",
880 "weakref",
881 "threading",
882 "multiprocessing",
883 "subprocess",
884 "queue",
885 "asyncio",
886 "socket",
887 "email",
888 "base64",
889 "hashlib",
890 "hmac",
891 "secrets",
892 "uuid",
893 "math",
894 "random",
895 "statistics",
896 "decimal",
897 "fractions",
898 "unittest",
899 "doctest",
900 "pdb",
901 "trace",
902 "timeit",
903 ];
904
905 let first_component = path.split('.').next().unwrap_or("");
907
908 STDLIB_MODULES.contains(&first_component)
909}
910
911fn classify_python_import(import_path: &str) -> ImportType {
913 if import_path.starts_with('.') {
915 return ImportType::Internal;
916 }
917
918 const STDLIB_MODULES: &[&str] = &[
920 "os",
921 "sys",
922 "io",
923 "re",
924 "json",
925 "csv",
926 "xml",
927 "html",
928 "http",
929 "urllib",
930 "collections",
931 "itertools",
932 "functools",
933 "operator",
934 "pathlib",
935 "glob",
936 "tempfile",
937 "shutil",
938 "pickle",
939 "shelve",
940 "sqlite3",
941 "zlib",
942 "gzip",
943 "time",
944 "datetime",
945 "calendar",
946 "logging",
947 "argparse",
948 "configparser",
949 "typing",
950 "dataclasses",
951 "enum",
952 "abc",
953 "contextlib",
954 "weakref",
955 "threading",
956 "multiprocessing",
957 "subprocess",
958 "queue",
959 "asyncio",
960 "socket",
961 "email",
962 "base64",
963 "hashlib",
964 "hmac",
965 "secrets",
966 "uuid",
967 "math",
968 "random",
969 "statistics",
970 "decimal",
971 "fractions",
972 "unittest",
973 "doctest",
974 "pdb",
975 "trace",
976 "timeit",
977 ];
978
979 let first_component = import_path.split('.').next().unwrap_or("");
981
982 if STDLIB_MODULES.contains(&first_component) {
983 ImportType::Stdlib
984 } else {
985 ImportType::External
987 }
988}
989
990#[derive(Debug, Clone)]
996pub struct PythonPackage {
997 pub name: String,
999 pub project_root: String,
1001 pub abs_project_root: std::path::PathBuf,
1003}
1004
1005pub fn find_all_python_configs(index_root: &std::path::Path) -> Result<Vec<std::path::PathBuf>> {
1008 use ignore::WalkBuilder;
1009
1010 let mut config_files = Vec::new();
1011
1012 let walker = WalkBuilder::new(index_root)
1013 .follow_links(false)
1014 .git_ignore(true)
1015 .build();
1016
1017 for entry in walker {
1018 let entry = entry?;
1019 let path = entry.path();
1020
1021 if !path.is_file() {
1022 continue;
1023 }
1024
1025 let filename = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
1026
1027 if filename == "pyproject.toml" || filename == "setup.py" || filename == "setup.cfg" {
1029 let path_str = path.to_string_lossy().replace('\\', "/");
1032 if path_str.contains("/venv/")
1033 || path_str.contains("/.venv/")
1034 || path_str.contains("/site-packages/")
1035 || path_str.contains("/dist/")
1036 || path_str.contains("/build/")
1037 || path_str.contains("/__pycache__/")
1038 {
1039 log::trace!(
1040 "Skipping Python config in vendor/build directory: {:?}",
1041 path
1042 );
1043 continue;
1044 }
1045
1046 config_files.push(path.to_path_buf());
1047 }
1048 }
1049
1050 log::debug!("Found {} Python config files", config_files.len());
1051 Ok(config_files)
1052}
1053
1054pub fn parse_all_python_packages(index_root: &std::path::Path) -> Result<Vec<PythonPackage>> {
1056 let config_files = find_all_python_configs(index_root)?;
1057
1058 if config_files.is_empty() {
1059 log::debug!("No Python config files found in {:?}", index_root);
1060 return Ok(Vec::new());
1061 }
1062
1063 let mut packages = Vec::new();
1064 let config_count = config_files.len();
1065
1066 for config_path in &config_files {
1067 let project_root = config_path
1068 .parent()
1069 .ok_or_else(|| anyhow::anyhow!("Config file has no parent directory"))?;
1070
1071 if let Some(package_name) = find_python_package_name(project_root) {
1073 let relative_project_root = project_root
1077 .strip_prefix(index_root)
1078 .unwrap_or(project_root)
1079 .to_string_lossy()
1080 .replace('\\', "/");
1081
1082 log::debug!(
1083 "Found Python package '{}' at {:?}",
1084 package_name,
1085 relative_project_root
1086 );
1087
1088 packages.push(PythonPackage {
1089 name: package_name,
1090 project_root: relative_project_root,
1091 abs_project_root: project_root.to_path_buf(),
1092 });
1093 }
1094 }
1095
1096 log::info!(
1097 "Loaded {} Python packages from {} config files",
1098 packages.len(),
1099 config_count
1100 );
1101
1102 Ok(packages)
1103}
1104
1105pub fn resolve_python_import_to_path(
1112 import_path: &str,
1113 packages: &[PythonPackage],
1114 current_file_path: Option<&str>,
1115) -> Option<String> {
1116 if import_path.starts_with('.') {
1118 return resolve_relative_python_import(import_path, current_file_path);
1119 }
1120
1121 let first_component = import_path.split('.').next()?;
1124
1125 for package in packages {
1127 if package.name == first_component {
1128 let module_path = import_path.replace('.', "/");
1131
1132 let candidates = vec![
1134 format!("{}/{}.py", package.project_root, module_path),
1135 format!("{}/{}/__init__.py", package.project_root, module_path),
1136 ];
1137
1138 if let Some(candidate) = candidates.into_iter().next() {
1139 log::trace!("Checking Python module path: {}", candidate);
1140 return Some(candidate);
1141 }
1142 }
1143 }
1144
1145 None
1146}
1147
1148fn resolve_relative_python_import(
1151 import_path: &str,
1152 current_file_path: Option<&str>,
1153) -> Option<String> {
1154 let current_file = current_file_path?;
1155
1156 let dots = import_path.chars().take_while(|&c| c == '.').count();
1158 if dots == 0 {
1159 return None;
1160 }
1161
1162 let current_dir = std::path::Path::new(current_file).parent()?;
1164
1165 let mut target_dir = current_dir.to_path_buf();
1167 for _ in 1..dots {
1168 target_dir = target_dir.parent()?.to_path_buf();
1169 }
1170
1171 let module_path = import_path.trim_start_matches('.');
1173
1174 if module_path.is_empty() {
1175 return Some(format!("{}/__init__.py", target_dir.to_string_lossy()));
1177 }
1178
1179 let file_path = module_path.replace('.', "/");
1181
1182 let candidates = vec![
1184 format!("{}/{}.py", target_dir.to_string_lossy(), file_path),
1185 format!("{}/{}/__init__.py", target_dir.to_string_lossy(), file_path),
1186 ];
1187
1188 if let Some(candidate) = candidates.into_iter().next() {
1189 log::trace!("Checking relative Python import: {}", candidate);
1190 return Some(candidate);
1191 }
1192
1193 None
1194}
1195
1196#[cfg(test)]
1197mod tests {
1198 use super::*;
1199
1200 #[test]
1201 fn test_parse_function() {
1202 let source = r#"
1203def hello_world():
1204 print("Hello, world!")
1205 return True
1206 "#;
1207
1208 let symbols = parse("test.py", source).unwrap();
1209 assert_eq!(symbols.len(), 1);
1210 assert_eq!(symbols[0].symbol.as_deref(), Some("hello_world"));
1211 assert!(matches!(symbols[0].kind, SymbolKind::Function));
1212 }
1213
1214 #[test]
1215 fn test_parse_async_function() {
1216 let source = r#"
1217async def fetch_data(url):
1218 async with aiohttp.ClientSession() as session:
1219 async with session.get(url) as response:
1220 return await response.text()
1221 "#;
1222
1223 let symbols = parse("test.py", source).unwrap();
1224 assert_eq!(symbols.len(), 1);
1225 assert_eq!(symbols[0].symbol.as_deref(), Some("fetch_data"));
1226 assert!(matches!(symbols[0].kind, SymbolKind::Function));
1227 }
1228
1229 #[test]
1230 fn test_parse_class() {
1231 let source = r#"
1232class User:
1233 def __init__(self, name, age):
1234 self.name = name
1235 self.age = age
1236 "#;
1237
1238 let symbols = parse("test.py", source).unwrap();
1239
1240 let class_symbols: Vec<_> = symbols
1241 .iter()
1242 .filter(|s| matches!(s.kind, SymbolKind::Class))
1243 .collect();
1244
1245 assert_eq!(class_symbols.len(), 1);
1246 assert_eq!(class_symbols[0].symbol.as_deref(), Some("User"));
1247 }
1248
1249 #[test]
1250 fn test_parse_class_with_methods() {
1251 let source = r#"
1252class Calculator:
1253 def add(self, a, b):
1254 return a + b
1255
1256 def subtract(self, a, b):
1257 return a - b
1258
1259 @staticmethod
1260 def multiply(a, b):
1261 return a * b
1262 "#;
1263
1264 let symbols = parse("test.py", source).unwrap();
1265
1266 let method_symbols: Vec<_> = symbols
1267 .iter()
1268 .filter(|s| matches!(s.kind, SymbolKind::Method))
1269 .collect();
1270
1271 assert_eq!(method_symbols.len(), 3);
1272 assert!(
1273 method_symbols
1274 .iter()
1275 .any(|s| s.symbol.as_deref() == Some("add"))
1276 );
1277 assert!(
1278 method_symbols
1279 .iter()
1280 .any(|s| s.symbol.as_deref() == Some("subtract"))
1281 );
1282 assert!(
1283 method_symbols
1284 .iter()
1285 .any(|s| s.symbol.as_deref() == Some("multiply"))
1286 );
1287
1288 for _method in method_symbols {
1290 }
1292 }
1293
1294 #[test]
1295 fn test_parse_async_method() {
1296 let source = r#"
1297class DataFetcher:
1298 async def get_user(self, user_id):
1299 return await fetch(f"/users/{user_id}")
1300
1301 async def get_all_users(self):
1302 return await fetch("/users")
1303 "#;
1304
1305 let symbols = parse("test.py", source).unwrap();
1306
1307 let method_symbols: Vec<_> = symbols
1308 .iter()
1309 .filter(|s| matches!(s.kind, SymbolKind::Method))
1310 .collect();
1311
1312 assert_eq!(method_symbols.len(), 2);
1313 assert!(
1314 method_symbols
1315 .iter()
1316 .any(|s| s.symbol.as_deref() == Some("get_user"))
1317 );
1318 assert!(
1319 method_symbols
1320 .iter()
1321 .any(|s| s.symbol.as_deref() == Some("get_all_users"))
1322 );
1323 }
1324
1325 #[test]
1326 fn test_parse_constants() {
1327 let source = r#"
1328MAX_SIZE = 100
1329DEFAULT_TIMEOUT = 30
1330API_URL = "https://api.example.com"
1331 "#;
1332
1333 let symbols = parse("test.py", source).unwrap();
1334
1335 let const_symbols: Vec<_> = symbols
1336 .iter()
1337 .filter(|s| matches!(s.kind, SymbolKind::Constant))
1338 .collect();
1339
1340 assert_eq!(const_symbols.len(), 3);
1341 assert!(
1342 const_symbols
1343 .iter()
1344 .any(|s| s.symbol.as_deref() == Some("MAX_SIZE"))
1345 );
1346 assert!(
1347 const_symbols
1348 .iter()
1349 .any(|s| s.symbol.as_deref() == Some("DEFAULT_TIMEOUT"))
1350 );
1351 assert!(
1352 const_symbols
1353 .iter()
1354 .any(|s| s.symbol.as_deref() == Some("API_URL"))
1355 );
1356 }
1357
1358 #[test]
1359 fn test_parse_lambda() {
1360 let source = r#"
1361square = lambda x: x * x
1362add = lambda a, b: a + b
1363 "#;
1364
1365 let symbols = parse("test.py", source).unwrap();
1366
1367 let lambda_symbols: Vec<_> = symbols
1368 .iter()
1369 .filter(|s| matches!(s.kind, SymbolKind::Function))
1370 .collect();
1371
1372 assert!(lambda_symbols.len() >= 2);
1373 assert!(
1374 lambda_symbols
1375 .iter()
1376 .any(|s| s.symbol.as_deref() == Some("square"))
1377 );
1378 assert!(
1379 lambda_symbols
1380 .iter()
1381 .any(|s| s.symbol.as_deref() == Some("add"))
1382 );
1383 }
1384
1385 #[test]
1386 fn test_parse_decorated_method() {
1387 let source = r#"
1388class WebService:
1389 @property
1390 def url(self):
1391 return self._url
1392
1393 @classmethod
1394 def from_config(cls, config):
1395 return cls(config['url'])
1396
1397 @staticmethod
1398 def validate_url(url):
1399 return url.startswith('http')
1400 "#;
1401
1402 let symbols = parse("test.py", source).unwrap();
1403
1404 let method_symbols: Vec<_> = symbols
1405 .iter()
1406 .filter(|s| matches!(s.kind, SymbolKind::Method))
1407 .collect();
1408
1409 assert_eq!(method_symbols.len(), 3);
1410 assert!(
1411 method_symbols
1412 .iter()
1413 .any(|s| s.symbol.as_deref() == Some("url"))
1414 );
1415 assert!(
1416 method_symbols
1417 .iter()
1418 .any(|s| s.symbol.as_deref() == Some("from_config"))
1419 );
1420 assert!(
1421 method_symbols
1422 .iter()
1423 .any(|s| s.symbol.as_deref() == Some("validate_url"))
1424 );
1425 }
1426
1427 #[test]
1428 fn test_parse_mixed_symbols() {
1429 let source = r#"
1430API_KEY = "secret123"
1431MAX_RETRIES = 3
1432
1433class APIClient:
1434 def __init__(self, api_key):
1435 self.api_key = api_key
1436
1437 async def request(self, endpoint):
1438 return await self._fetch(endpoint)
1439
1440 @staticmethod
1441 def build_url(endpoint):
1442 return f"https://api.example.com/{endpoint}"
1443
1444def create_client():
1445 return APIClient(API_KEY)
1446
1447process = lambda data: data.strip().lower()
1448 "#;
1449
1450 let symbols = parse("test.py", source).unwrap();
1451
1452 assert!(symbols.len() >= 8);
1454
1455 let kinds: Vec<&SymbolKind> = symbols.iter().map(|s| &s.kind).collect();
1456 assert!(kinds.contains(&&SymbolKind::Constant));
1457 assert!(kinds.contains(&&SymbolKind::Class));
1458 assert!(kinds.contains(&&SymbolKind::Method));
1459 assert!(kinds.contains(&&SymbolKind::Function));
1460 }
1461
1462 #[test]
1463 fn test_parse_nested_class() {
1464 let source = r#"
1465class Outer:
1466 class Inner:
1467 def inner_method(self):
1468 pass
1469
1470 def outer_method(self):
1471 pass
1472 "#;
1473
1474 let symbols = parse("test.py", source).unwrap();
1475
1476 let class_symbols: Vec<_> = symbols
1477 .iter()
1478 .filter(|s| matches!(s.kind, SymbolKind::Class))
1479 .collect();
1480
1481 assert_eq!(class_symbols.len(), 2);
1483 assert!(
1484 class_symbols
1485 .iter()
1486 .any(|s| s.symbol.as_deref() == Some("Outer"))
1487 );
1488 assert!(
1489 class_symbols
1490 .iter()
1491 .any(|s| s.symbol.as_deref() == Some("Inner"))
1492 );
1493 }
1494
1495 #[test]
1496 fn test_local_variables_included() {
1497 let source = r#"
1498def calculate(input):
1499 local_var = input * 2
1500 result = local_var + 10
1501 return result
1502
1503class Calculator:
1504 def compute(self, value):
1505 temp = value * 3
1506 final = temp + 5
1507 return final
1508 "#;
1509
1510 let symbols = parse("test.py", source).unwrap();
1511
1512 let variables: Vec<_> = symbols
1514 .iter()
1515 .filter(|s| matches!(s.kind, SymbolKind::Variable))
1516 .collect();
1517
1518 assert!(
1520 variables
1521 .iter()
1522 .any(|v| v.symbol.as_deref() == Some("local_var"))
1523 );
1524 assert!(
1525 variables
1526 .iter()
1527 .any(|v| v.symbol.as_deref() == Some("result"))
1528 );
1529 assert!(
1530 variables
1531 .iter()
1532 .any(|v| v.symbol.as_deref() == Some("temp"))
1533 );
1534 assert!(
1535 variables
1536 .iter()
1537 .any(|v| v.symbol.as_deref() == Some("final"))
1538 );
1539
1540 for _var in variables {
1542 }
1544 }
1545
1546 #[test]
1547 fn test_global_variables() {
1548 let source = r#"
1549# Global constants (uppercase)
1550MAX_SIZE = 100
1551DEFAULT_TIMEOUT = 30
1552
1553# Global variables (non-uppercase)
1554database_url = "postgresql://localhost/mydb"
1555config = {"debug": True}
1556current_user = None
1557
1558def get_config():
1559 return config
1560 "#;
1561
1562 let symbols = parse("test.py", source).unwrap();
1563
1564 let constants: Vec<_> = symbols
1566 .iter()
1567 .filter(|s| matches!(s.kind, SymbolKind::Constant))
1568 .collect();
1569
1570 let variables: Vec<_> = symbols
1571 .iter()
1572 .filter(|s| matches!(s.kind, SymbolKind::Variable))
1573 .collect();
1574
1575 assert!(
1577 constants
1578 .iter()
1579 .any(|c| c.symbol.as_deref() == Some("MAX_SIZE"))
1580 );
1581 assert!(
1582 constants
1583 .iter()
1584 .any(|c| c.symbol.as_deref() == Some("DEFAULT_TIMEOUT"))
1585 );
1586
1587 assert!(
1589 variables
1590 .iter()
1591 .any(|v| v.symbol.as_deref() == Some("database_url"))
1592 );
1593 assert!(
1594 variables
1595 .iter()
1596 .any(|v| v.symbol.as_deref() == Some("config"))
1597 );
1598 assert!(
1599 variables
1600 .iter()
1601 .any(|v| v.symbol.as_deref() == Some("current_user"))
1602 );
1603
1604 for _constant in constants {
1606 }
1608 for _var in variables {
1609 }
1611 }
1612
1613 #[test]
1614 fn test_find_all_python_configs() {
1615 use std::fs;
1616 use tempfile::TempDir;
1617
1618 let temp = TempDir::new().unwrap();
1619 let root = temp.path();
1620
1621 let project1 = root.join("backend");
1623 fs::create_dir_all(&project1).unwrap();
1624 fs::write(
1625 project1.join("pyproject.toml"),
1626 "[project]\nname = \"backend\"",
1627 )
1628 .unwrap();
1629
1630 let project2 = root.join("frontend/api");
1631 fs::create_dir_all(&project2).unwrap();
1632 fs::write(project2.join("setup.py"), "setup(name='api')").unwrap();
1633
1634 let venv = root.join("venv");
1636 fs::create_dir_all(&venv).unwrap();
1637 fs::write(venv.join("setup.py"), "setup(name='should_skip')").unwrap();
1638
1639 let configs = find_all_python_configs(root).unwrap();
1640
1641 assert_eq!(configs.len(), 2);
1643 assert!(
1644 configs
1645 .iter()
1646 .any(|p| p.ends_with("backend/pyproject.toml"))
1647 );
1648 assert!(configs.iter().any(|p| p.ends_with("frontend/api/setup.py")));
1649 }
1650
1651 #[test]
1652 fn test_parse_all_python_packages() {
1653 use std::fs;
1654 use tempfile::TempDir;
1655
1656 let temp = TempDir::new().unwrap();
1657 let root = temp.path();
1658
1659 let project1 = root.join("services/auth");
1661 fs::create_dir_all(&project1).unwrap();
1662 fs::write(
1663 project1.join("pyproject.toml"),
1664 "[project]\nname = \"auth-service\"\n",
1665 )
1666 .unwrap();
1667
1668 let project2 = root.join("services/api");
1669 fs::create_dir_all(&project2).unwrap();
1670 fs::write(project2.join("setup.py"), "setup(name=\"api-service\")").unwrap();
1671
1672 let packages = parse_all_python_packages(root).unwrap();
1673
1674 assert_eq!(packages.len(), 2);
1676
1677 let names: Vec<_> = packages.iter().map(|p| p.name.as_str()).collect();
1679 assert!(names.contains(&"auth-service"));
1680 assert!(names.contains(&"api-service"));
1681
1682 for package in &packages {
1684 assert!(package.project_root.starts_with("services/"));
1685 assert!(package.abs_project_root.ends_with(&package.project_root));
1686 }
1687 }
1688
1689 #[test]
1690 fn test_resolve_python_import_absolute() {
1691 use std::fs;
1692 use tempfile::TempDir;
1693
1694 let temp = TempDir::new().unwrap();
1695 let root = temp.path();
1696
1697 let myapp = root.join("myapp");
1699 fs::create_dir_all(myapp.join("models")).unwrap();
1700 fs::write(
1701 myapp.join("pyproject.toml"),
1702 "[project]\nname = \"myapp\"\n",
1703 )
1704 .unwrap();
1705
1706 let packages = parse_all_python_packages(root).unwrap();
1707 assert_eq!(packages.len(), 1);
1708
1709 let resolved = resolve_python_import_to_path("myapp.models.user", &packages, None);
1712
1713 assert!(resolved.is_some());
1714 let path = resolved.unwrap();
1715 assert!(
1716 path.contains("myapp/models/user.py") || path.contains("myapp/models/user/__init__.py")
1717 );
1718 }
1719
1720 #[test]
1721 fn test_resolve_python_import_relative() {
1722 let current_file = "myapp/views/admin.py";
1724
1725 let resolved = resolve_python_import_to_path(
1727 ".models",
1728 &[], Some(current_file),
1730 );
1731
1732 assert!(resolved.is_some());
1733 let path = resolved.unwrap();
1734 assert!(path.contains("myapp/views/models"));
1736
1737 let resolved = resolve_python_import_to_path("..utils", &[], Some(current_file));
1739
1740 assert!(resolved.is_some());
1741 let path = resolved.unwrap();
1742 assert!(path.contains("myapp/utils"));
1744 }
1745
1746 #[test]
1747 fn test_resolve_python_import_relative_with_module() {
1748 let current_file = "myapp/views/dashboard/index.py";
1750
1751 let resolved = resolve_python_import_to_path("..models.user", &[], Some(current_file));
1752
1753 assert!(resolved.is_some());
1754 let path = resolved.unwrap();
1755 assert!(path.contains("models/user"));
1757 }
1758
1759 #[test]
1760 fn test_resolve_python_import_not_found() {
1761 use std::fs;
1762 use tempfile::TempDir;
1763
1764 let temp = TempDir::new().unwrap();
1765 let root = temp.path();
1766
1767 let myapp = root.join("myapp");
1768 fs::create_dir_all(&myapp).unwrap();
1769 fs::write(
1770 myapp.join("pyproject.toml"),
1771 "[project]\nname = \"myapp\"\n",
1772 )
1773 .unwrap();
1774
1775 let packages = parse_all_python_packages(root).unwrap();
1776
1777 let resolved = resolve_python_import_to_path("other_package.module", &packages, None);
1779
1780 assert!(resolved.is_none());
1782 }
1783
1784 #[test]
1785 fn test_dynamic_imports_filtered() {
1786 let source = r#"
1787import os
1788import sys
1789from json import loads
1790from .models import User
1791
1792# Dynamic imports - should be filtered out
1793import importlib
1794mod = importlib.import_module("some_module")
1795pkg = __import__("package")
1796exec("import dynamic")
1797 "#;
1798
1799 let deps = PythonDependencyExtractor::extract_dependencies(source).unwrap();
1800
1801 assert_eq!(deps.len(), 5, "Should extract 5 static imports only");
1804
1805 assert!(deps.iter().any(|d| d.imported_path == "os"));
1806 assert!(deps.iter().any(|d| d.imported_path == "sys"));
1807 assert!(deps.iter().any(|d| d.imported_path == "json"));
1808 assert!(deps.iter().any(|d| d.imported_path == ".models"));
1809 assert!(deps.iter().any(|d| d.imported_path == "importlib"));
1810
1811 assert!(!deps.iter().any(|d| d.imported_path.contains("some_module")));
1813 assert!(
1814 !deps
1815 .iter()
1816 .any(|d| d.imported_path.contains("package") && d.imported_path != "json")
1817 );
1818 assert!(!deps.iter().any(|d| d.imported_path.contains("dynamic")));
1819 }
1820}