1use std::{
2 collections::HashMap,
3 path::{Path, PathBuf},
4};
5
6use gitcortex_core::{
7 error::{GitCortexError, Result},
8 graph::{Edge, Node, NodeId, NodeMetadata, Span},
9 schema::{EdgeConfidence, EdgeKind, NodeKind, Visibility},
10};
11use tree_sitter::{Node as TsNode, Parser};
12
13use super::{capture_definition, LanguageParser, ParseResult};
14
15pub struct PythonParser {
16 language: tree_sitter::Language,
17}
18
19impl PythonParser {
20 pub fn new() -> Self {
21 Self {
22 language: tree_sitter_python::LANGUAGE.into(),
23 }
24 }
25}
26
27impl Default for PythonParser {
28 fn default() -> Self {
29 Self::new()
30 }
31}
32
33impl LanguageParser for PythonParser {
34 fn extensions(&self) -> &[&str] {
35 &["py"]
36 }
37
38 fn parse(&self, path: &Path, source: &str) -> Result<ParseResult> {
39 let mut parser = Parser::new();
40 parser
41 .set_language(&self.language)
42 .map_err(|e| GitCortexError::Parse {
43 file: path.to_owned(),
44 message: e.to_string(),
45 })?;
46
47 let tree = parser
48 .parse(source, None)
49 .ok_or_else(|| GitCortexError::Parse {
50 file: path.to_owned(),
51 message: "tree-sitter returned no parse tree".into(),
52 })?;
53
54 let mut visitor = FileVisitor::new(path, source);
55 visitor.collect_names(tree.root_node());
56 visitor.visit_module(tree.root_node());
57 visitor.collect_imports(tree.root_node());
58
59 Ok(ParseResult {
60 nodes: visitor.nodes,
61 edges: visitor.edges,
62 deferred_calls: visitor.deferred_calls,
63 deferred_uses: visitor.deferred_uses,
64 deferred_implements: visitor.deferred_implements,
65 deferred_imports: visitor.deferred_imports,
66 deferred_inherits: Vec::new(),
67 deferred_throws: Vec::new(),
68 deferred_annotated: visitor.deferred_annotated,
69 deferred_doc_refs: Vec::new(),
70 })
71 }
72}
73
74struct FileVisitor<'src> {
77 source: &'src [u8],
78 file: PathBuf,
79 module_id: NodeId,
81 nodes: Vec<Node>,
82 edges: Vec<Edge>,
83 class_index: HashMap<String, NodeId>,
85 fn_index: HashMap<String, NodeId>,
87 deferred_calls: Vec<(NodeId, String, u32)>,
88 deferred_uses: Vec<(NodeId, String)>,
89 deferred_implements: Vec<(NodeId, String)>,
90 deferred_imports: Vec<(NodeId, String)>,
91 deferred_annotated: Vec<(NodeId, String)>,
92}
93
94impl<'src> FileVisitor<'src> {
95 fn new(file: &Path, source: &'src str) -> Self {
96 let module_id = NodeId::new();
97 let module_name = file
99 .file_stem()
100 .and_then(|s| s.to_str())
101 .unwrap_or("__init__")
102 .to_owned();
103 let module_node = Node {
104 id: module_id.clone(),
105 qualified_name: module_name.clone(),
106 kind: NodeKind::Module,
107 name: module_name,
108 file: file.to_owned(),
109 span: Span {
110 start_line: 1,
111 end_line: 1,
112 },
113 metadata: NodeMetadata {
114 loc: source.lines().count() as u32,
115 visibility: Visibility::Pub,
116 is_async: false,
117 is_unsafe: false,
118 ..Default::default()
119 },
120 };
121 let nodes = vec![module_node];
122 Self {
123 source: source.as_bytes(),
124 file: file.to_owned(),
125 module_id,
126 nodes,
127 edges: Vec::new(),
128 class_index: HashMap::new(),
129 fn_index: HashMap::new(),
130 deferred_calls: Vec::new(),
131 deferred_uses: Vec::new(),
132 deferred_implements: Vec::new(),
133 deferred_imports: Vec::new(),
134 deferred_annotated: Vec::new(),
135 }
136 }
137
138 fn text<'t>(&self, node: TsNode<'t>) -> &'src str {
139 node.utf8_text(self.source).unwrap_or("")
140 }
141
142 fn span(node: TsNode<'_>) -> Span {
143 Span {
144 start_line: node.start_position().row as u32 + 1,
145 end_line: node.end_position().row as u32 + 1,
146 }
147 }
148
149 fn visibility(name: &str) -> Visibility {
151 if name.starts_with('_') {
152 Visibility::Private
153 } else {
154 Visibility::Pub
155 }
156 }
157
158 fn qualified(scope: &[String], name: &str) -> String {
159 if scope.is_empty() {
160 name.to_owned()
161 } else {
162 format!("{}.{name}", scope.join("."))
163 }
164 }
165
166 fn make_node(
167 &self,
168 id: NodeId,
169 kind: NodeKind,
170 name: String,
171 scope: &[String],
172 ts_node: TsNode<'_>,
173 is_async: bool,
174 ) -> Node {
175 let vis = Self::visibility(&name);
176 Node {
177 id,
178 qualified_name: Self::qualified(scope, &name),
179 kind,
180 name,
181 file: self.file.clone(),
182 span: Self::span(ts_node),
183 metadata: NodeMetadata {
184 loc: (ts_node.end_position().row - ts_node.start_position().row + 1) as u32,
185 visibility: vis,
186 is_async,
187 is_unsafe: false,
188 definition: capture_definition(self.source, ts_node),
189 ..Default::default()
190 },
191 }
192 }
193
194 fn collect_names(&mut self, node: TsNode<'_>) {
197 let mut cursor = node.walk();
198 let children: Vec<TsNode<'_>> = node.named_children(&mut cursor).collect();
199 for child in children {
200 match child.kind() {
201 "class_definition" => {
202 if let Some(name_node) = child.child_by_field_name("name") {
203 let name = self.text(name_node).to_owned();
204 self.class_index.entry(name).or_default();
205 }
206 if let Some(body) = child.child_by_field_name("body") {
208 self.collect_names(body);
209 }
210 }
211 "function_definition" => {
212 if let Some(name_node) = child.child_by_field_name("name") {
213 let name = self.text(name_node).to_owned();
214 self.fn_index.entry(name).or_default();
215 }
216 }
217 "decorated_definition" => {
218 let def = child.child_by_field_name("definition");
219 if let Some(def) = def {
220 match def.kind() {
221 "function_definition" => {
222 if let Some(name_node) = def.child_by_field_name("name") {
223 let name = self.text(name_node).to_owned();
224 self.fn_index.entry(name).or_default();
225 }
226 }
227 "class_definition" => {
228 if let Some(name_node) = def.child_by_field_name("name") {
229 let name = self.text(name_node).to_owned();
230 self.class_index.entry(name).or_default();
231 }
232 if let Some(body) = def.child_by_field_name("body") {
234 self.collect_names(body);
235 }
236 }
237 _ => {}
238 }
239 }
240 }
241 _ => {}
242 }
243 }
244 }
245
246 fn visit_module(&mut self, node: TsNode<'_>) {
249 let mut cursor = node.walk();
250 let children: Vec<TsNode<'_>> = node.named_children(&mut cursor).collect();
251 for child in children {
252 self.visit_top_level(child, &[]);
253 }
254 }
255
256 fn visit_top_level(&mut self, node: TsNode<'_>, scope: &[String]) {
257 match node.kind() {
258 "function_definition" => {
259 let is_async = Self::fn_is_async(node);
260 self.visit_function(node, scope, None, is_async, &[]);
261 }
262 "decorated_definition" => {
263 let decorators = self.collect_decorators(node);
264 let is_async = node
265 .child_by_field_name("definition")
266 .map(Self::fn_is_async)
267 .unwrap_or(false);
268 if let Some(def) = node.child_by_field_name("definition") {
269 match def.kind() {
270 "function_definition" => {
271 self.visit_function(def, scope, None, is_async, &decorators)
272 }
273 "class_definition" => self.visit_class(def, scope, &decorators),
274 _ => {}
275 }
276 }
277 }
278 "class_definition" => self.visit_class(node, scope, &[]),
279 "expression_statement" => self.maybe_visit_constant(node, scope),
280 _ => {}
281 }
282 }
283
284 fn visit_function(
285 &mut self,
286 node: TsNode<'_>,
287 scope: &[String],
288 container_id: Option<NodeId>,
289 is_async: bool,
290 decorators: &[String],
291 ) {
292 let Some(name_node) = node.child_by_field_name("name") else {
293 return;
294 };
295 let name = self.text(name_node).to_owned();
296 let id = self
297 .fn_index
298 .get(&name)
299 .cloned()
300 .unwrap_or_else(NodeId::new);
301
302 let has_property = decorators.iter().any(|d| d == "property");
304 let has_staticmethod = decorators.iter().any(|d| d == "staticmethod");
305 let has_classmethod = decorators.iter().any(|d| d == "classmethod");
306
307 let kind = if has_property {
308 NodeKind::Property
309 } else if container_id.is_some() {
310 NodeKind::Method
311 } else {
312 NodeKind::Function
313 };
314
315 let is_generator = node
317 .child_by_field_name("body")
318 .map(|body| Self::body_has_yield(body))
319 .unwrap_or(false);
320
321 let mut graph_node = self.make_node(id.clone(), kind, name, scope, node, is_async);
322 if has_property {
323 graph_node.metadata.is_property = true;
324 }
325 if has_staticmethod || has_classmethod {
326 graph_node.metadata.is_static = true;
327 }
328 if is_generator {
329 graph_node.metadata.is_generator = true;
330 }
331
332 if let Some(body) = node.child_by_field_name("body") {
333 graph_node.metadata.lld.complexity = Some(super::cyclomatic_complexity(
334 body,
335 &super::complexity::python_decision,
336 ));
337 }
338
339 if let Some(cid) = container_id {
340 self.edges.push(Edge {
341 src: cid,
342 dst: id.clone(),
343 kind: EdgeKind::Contains,
344 line: None,
345 confidence: EdgeConfidence::Extracted,
346 });
347 }
348 self.nodes.push(graph_node);
349
350 self.extract_param_types(node, &id);
352 self.extract_return_type(node, &id);
353
354 for dec in decorators {
357 self.deferred_uses.push((id.clone(), dec.clone()));
358 self.deferred_annotated.push((id.clone(), dec.clone()));
359 }
360
361 if let Some(body) = node.child_by_field_name("body") {
362 self.collect_calls(body, &id);
363 }
364 }
365
366 fn visit_class(&mut self, node: TsNode<'_>, scope: &[String], decorators: &[String]) {
367 let Some(name_node) = node.child_by_field_name("name") else {
368 return;
369 };
370 let name = self.text(name_node).to_owned();
371 let id = self
372 .class_index
373 .get(&name)
374 .cloned()
375 .unwrap_or_else(NodeId::new);
376
377 let mut is_protocol = false;
379 if let Some(bases) = node.child_by_field_name("superclasses") {
380 let mut c = bases.walk();
381 for base in bases.named_children(&mut c) {
382 let base_name = match base.kind() {
383 "identifier" => Some(self.text(base).to_owned()),
384 "attribute" => base
385 .child_by_field_name("attribute")
386 .map(|n| self.text(n).to_owned()),
387 _ => None,
388 };
389 if base_name.as_deref() == Some("Protocol") {
390 is_protocol = true;
391 }
392 }
393 }
394
395 let class_kind = if is_protocol {
396 NodeKind::Interface
397 } else {
398 NodeKind::Struct
399 };
400
401 let mut graph_node =
402 self.make_node(id.clone(), class_kind, name.clone(), scope, node, false);
403 if is_protocol {
404 graph_node.metadata.is_abstract = true;
405 }
406 self.nodes.push(graph_node);
407
408 if let Some(bases) = node.child_by_field_name("superclasses") {
410 let mut c = bases.walk();
411 for base in bases.named_children(&mut c) {
412 let base_name = match base.kind() {
413 "identifier" => Some(self.text(base).to_owned()),
414 "attribute" => base
415 .child_by_field_name("attribute")
416 .map(|n| self.text(n).to_owned()),
417 _ => None,
418 };
419 if let Some(b) = base_name {
420 self.deferred_implements.push((id.clone(), b));
421 }
422 }
423 }
424
425 for dec in decorators {
427 self.deferred_uses.push((id.clone(), dec.clone()));
428 self.deferred_annotated.push((id.clone(), dec.clone()));
429 }
430
431 let mut class_scope = scope.to_vec();
432 class_scope.push(name.clone());
433
434 if let Some(body) = node.child_by_field_name("body") {
435 let mut cursor = body.walk();
436 let children: Vec<TsNode<'_>> = body.named_children(&mut cursor).collect();
437 for child in children {
438 match child.kind() {
439 "function_definition" => {
440 let is_async = Self::fn_is_async(child);
441 self.visit_function(child, &class_scope, Some(id.clone()), is_async, &[]);
442 }
443 "decorated_definition" => {
444 let method_decorators = self.collect_decorators(child);
445 let is_async = child
446 .child_by_field_name("definition")
447 .map(Self::fn_is_async)
448 .unwrap_or(false);
449 if let Some(def) = child.child_by_field_name("definition") {
450 match def.kind() {
451 "function_definition" => {
452 self.visit_function(
453 def,
454 &class_scope,
455 Some(id.clone()),
456 is_async,
457 &method_decorators,
458 );
459 }
460 "class_definition" => {
461 self.visit_class(def, &class_scope, &method_decorators);
462 if let Some(nested_name_node) = def.child_by_field_name("name")
464 {
465 let nested_name = self.text(nested_name_node).to_owned();
466 if let Some(nested_id) =
467 self.class_index.get(&nested_name).cloned()
468 {
469 self.edges.push(Edge {
470 src: id.clone(),
471 dst: nested_id,
472 kind: EdgeKind::Contains,
473 line: None,
474 confidence: EdgeConfidence::Extracted,
475 });
476 }
477 }
478 }
479 _ => {}
480 }
481 }
482 }
483 "class_definition" => {
484 self.visit_class(child, &class_scope, &[]);
485 if let Some(nested_name_node) = child.child_by_field_name("name") {
487 let nested_name = self.text(nested_name_node).to_owned();
488 if let Some(nested_id) = self.class_index.get(&nested_name).cloned() {
489 self.edges.push(Edge {
490 src: id.clone(),
491 dst: nested_id,
492 kind: EdgeKind::Contains,
493 line: None,
494 confidence: EdgeConfidence::Extracted,
495 });
496 }
497 }
498 }
499 _ => {}
500 }
501 }
502 }
503 }
504
505 fn maybe_visit_constant(&mut self, node: TsNode<'_>, scope: &[String]) {
506 let mut cursor = node.walk();
517 for child in node.named_children(&mut cursor) {
518 if child.kind() == "assignment" {
519 if let Some(left) = child.child_by_field_name("left") {
520 if left.kind() == "identifier" {
521 let name = self.text(left).to_owned();
522 if name.is_empty() {
523 continue;
524 }
525 let id = NodeId::new();
526 let graph_node =
527 self.make_node(id, NodeKind::Constant, name, scope, node, false);
528 self.nodes.push(graph_node);
529 }
530 }
531 }
532 }
533 }
534
535 fn collect_imports(&mut self, node: TsNode<'_>) {
538 let mut cursor = node.walk();
539 let children: Vec<TsNode<'_>> = node.named_children(&mut cursor).collect();
540 for child in children {
541 match child.kind() {
542 "import_statement" => {
543 let mut c = child.walk();
545 for name_node in child.named_children(&mut c) {
546 let leaf = match name_node.kind() {
547 "dotted_name" => {
548 let text = self.text(name_node);
549 text.split('.').next_back().map(|s| s.to_owned())
550 }
551 "aliased_import" => name_node
552 .child_by_field_name("name")
553 .map(|n| self.text(n))
554 .map(|t| t.split('.').next_back().unwrap_or(t).to_owned()),
555 _ => None,
556 };
557 if let Some(name) = leaf {
558 self.deferred_imports.push((self.module_id.clone(), name));
559 }
560 }
561 }
562 "import_from_statement" => {
563 let mut c = child.walk();
567 let all_children: Vec<TsNode<'_>> = child.named_children(&mut c).collect();
568 for name_node in all_children.iter().skip(1) {
569 let leaf = match name_node.kind() {
570 "dotted_name" => {
571 let text = self.text(*name_node);
572 Some(text.split('.').next_back().unwrap_or(text).to_owned())
573 }
574 "aliased_import" => name_node
575 .child_by_field_name("name")
576 .map(|n| self.text(n))
577 .map(|t| t.split('.').next_back().unwrap_or(t).to_owned()),
578 "wildcard_import" => None,
579 _ => None,
580 };
581 if let Some(name) = leaf {
582 self.deferred_imports.push((self.module_id.clone(), name));
583 }
584 }
585 }
586 _ => {}
587 }
588 }
589 }
590
591 fn collect_calls(&mut self, node: TsNode<'_>, caller_id: &NodeId) {
594 let mut cursor = node.walk();
595 let children: Vec<TsNode<'_>> = node.named_children(&mut cursor).collect();
596 for child in children {
597 if child.kind() == "call" {
598 if let Some(callee) = self.callee_name(child) {
599 let line = child.start_position().row as u32 + 1;
600 self.record_call(caller_id.clone(), callee, line);
601 }
602 if let Some(args) = child.child_by_field_name("arguments") {
603 self.collect_calls(args, caller_id);
604 }
605 } else {
606 self.collect_calls(child, caller_id);
607 }
608 }
609 }
610
611 fn callee_name(&self, call_node: TsNode<'_>) -> Option<String> {
612 let func = call_node.child_by_field_name("function")?;
613 match func.kind() {
614 "identifier" => Some(self.text(func).to_owned()),
615 "attribute" => func
616 .child_by_field_name("attribute")
617 .map(|n| self.text(n).to_owned()),
618 _ => None,
619 }
620 }
621
622 fn record_call(&mut self, caller_id: NodeId, callee_name: String, line: u32) {
623 if callee_name.is_empty() {
624 return;
625 }
626 if let Some(callee_id) = self.fn_index.get(&callee_name).cloned() {
627 let edge = Edge::call(caller_id, callee_id, line);
628 if !self.edges.contains(&edge) {
629 self.edges.push(edge);
630 }
631 } else if !self
632 .deferred_calls
633 .iter()
634 .any(|(c, n, _)| c == &caller_id && n == &callee_name)
635 {
636 self.deferred_calls.push((caller_id, callee_name, line));
637 }
638 }
639
640 fn fn_is_async(node: TsNode<'_>) -> bool {
644 let mut c = node.walk();
645 let result = node.children(&mut c).any(|n| n.kind() == "async");
647 result
648 }
649
650 fn body_has_yield(node: TsNode<'_>) -> bool {
652 if node.kind() == "yield" || node.kind() == "yield_from" {
653 return true;
654 }
655 if node.kind() == "function_definition" {
658 return false;
659 }
660 let mut c = node.walk();
661 let found = node.named_children(&mut c).any(Self::body_has_yield);
662 found
663 }
664
665 fn collect_decorators(&self, node: TsNode<'_>) -> Vec<String> {
667 let mut c = node.walk();
668 node.named_children(&mut c)
669 .filter(|n| n.kind() == "decorator")
670 .filter_map(|d| self.decorator_name(d))
671 .collect()
672 }
673
674 fn decorator_name(&self, decorator: TsNode<'_>) -> Option<String> {
676 let mut c = decorator.walk();
677 let child = decorator.named_children(&mut c).next()?;
678 match child.kind() {
679 "identifier" => Some(self.text(child).to_owned()),
680 "attribute" => child
681 .child_by_field_name("attribute")
682 .map(|n| self.text(n).to_owned()),
683 "call" => child
684 .child_by_field_name("function")
685 .and_then(|f| match f.kind() {
686 "identifier" => Some(self.text(f).to_owned()),
687 "attribute" => f
688 .child_by_field_name("attribute")
689 .map(|n| self.text(n).to_owned()),
690 _ => None,
691 }),
692 _ => None,
693 }
694 }
695
696 fn extract_param_types(&mut self, fn_node: TsNode<'_>, fn_id: &NodeId) {
699 let Some(params) = fn_node.child_by_field_name("parameters") else {
700 return;
701 };
702 let mut c = params.walk();
703 for param in params.named_children(&mut c) {
704 let type_node = match param.kind() {
705 "typed_parameter" | "typed_default_parameter" => param.child_by_field_name("type"),
706 _ => None,
707 };
708 if let Some(t) = type_node {
709 for name in self.collect_type_names(t) {
710 self.deferred_uses.push((fn_id.clone(), name));
711 }
712 }
713 }
714 }
715
716 fn extract_return_type(&mut self, fn_node: TsNode<'_>, fn_id: &NodeId) {
717 if let Some(ret) = fn_node.child_by_field_name("return_type") {
718 for name in self.collect_type_names(ret) {
719 self.deferred_uses.push((fn_id.clone(), name));
720 }
721 }
722 }
723
724 fn collect_type_names(&self, node: TsNode<'_>) -> Vec<String> {
726 let mut names = Vec::new();
727 self.walk_type_names(node, &mut names);
728 names
729 }
730
731 fn walk_type_names(&self, node: TsNode<'_>, out: &mut Vec<String>) {
732 match node.kind() {
733 "identifier" => {
734 let name = self.text(node).to_owned();
735 if !is_builtin_type(&name) {
736 out.push(name);
737 }
738 }
739 _ => {
740 let mut c = node.walk();
741 for child in node.named_children(&mut c) {
742 self.walk_type_names(child, out);
743 }
744 }
745 }
746 }
747}
748
749fn is_builtin_type(name: &str) -> bool {
752 matches!(
753 name,
754 "int"
755 | "str"
756 | "bool"
757 | "float"
758 | "complex"
759 | "bytes"
760 | "bytearray"
761 | "None"
762 | "list"
763 | "dict"
764 | "set"
765 | "frozenset"
766 | "tuple"
767 | "type"
768 | "object"
769 | "Any"
770 | "Optional"
771 | "Union"
772 | "List"
773 | "Dict"
774 | "Set"
775 | "FrozenSet"
776 | "Tuple"
777 | "Callable"
778 | "Type"
779 | "ClassVar"
780 | "Final"
781 | "Literal"
782 | "TypeVar"
783 | "Generic"
784 | "Protocol"
785 | "Sequence"
786 | "Iterable"
787 | "Iterator"
788 | "Generator"
789 | "Coroutine"
790 | "Awaitable"
791 | "AsyncIterator"
792 | "AsyncGenerator"
793 | "NoReturn"
794 | "Never"
795 | "Self"
796 | "Annotated"
797 | "TypeAlias"
798 | "ParamSpec"
799 | "TypeVarTuple"
800 | "overload"
801 | "abstractmethod"
802 | "staticmethod"
803 | "classmethod"
804 )
805}
806
807#[cfg(test)]
810mod tests {
811 use super::PythonParser;
812 use crate::parser::LanguageParser;
813 use gitcortex_core::schema::{EdgeKind, NodeKind};
814 use std::path::Path;
815
816 fn parse(
817 src: &str,
818 ) -> (
819 Vec<gitcortex_core::graph::Node>,
820 Vec<gitcortex_core::graph::Edge>,
821 ) {
822 let r = PythonParser::new()
823 .parse(Path::new("test.py"), src)
824 .unwrap();
825 (r.nodes, r.edges)
826 }
827
828 #[allow(clippy::type_complexity)]
829 fn parse_full(
830 src: &str,
831 ) -> (
832 Vec<gitcortex_core::graph::Node>,
833 Vec<gitcortex_core::graph::Edge>,
834 Vec<(gitcortex_core::graph::NodeId, String, u32)>,
835 Vec<(gitcortex_core::graph::NodeId, String)>,
836 Vec<(gitcortex_core::graph::NodeId, String)>,
837 Vec<(gitcortex_core::graph::NodeId, String)>,
838 ) {
839 let r = PythonParser::new()
840 .parse(Path::new("test.py"), src)
841 .unwrap();
842 (
843 r.nodes,
844 r.edges,
845 r.deferred_calls,
846 r.deferred_uses,
847 r.deferred_implements,
848 r.deferred_imports,
849 )
850 }
851
852 #[test]
853 fn parses_free_function() {
854 let (nodes, _) = parse("def greet(name):\n return name\n");
855 assert_eq!(nodes.len(), 2);
857 let fns: Vec<_> = nodes
858 .iter()
859 .filter(|n| n.kind == NodeKind::Function)
860 .collect();
861 assert_eq!(fns.len(), 1);
862 assert_eq!(fns[0].name, "greet");
863 }
864
865 #[test]
866 fn parses_class_and_method() {
867 let src = "class Person:\n def greet(self):\n pass\n";
868 let (nodes, edges) = parse(src);
869 let classes: Vec<_> = nodes
870 .iter()
871 .filter(|n| n.kind == NodeKind::Struct)
872 .collect();
873 let methods: Vec<_> = nodes
874 .iter()
875 .filter(|n| n.kind == NodeKind::Method)
876 .collect();
877 assert_eq!(classes.len(), 1);
878 assert_eq!(methods.len(), 1);
879 let contains: Vec<_> = edges
880 .iter()
881 .filter(|e| e.kind == EdgeKind::Contains)
882 .collect();
883 assert!(!contains.is_empty());
884 }
885
886 #[test]
887 fn detects_call_edges() {
888 let src = "def caller():\n callee()\ndef callee():\n pass\n";
889 let (_, edges) = parse(src);
890 let calls: Vec<_> = edges.iter().filter(|e| e.kind == EdgeKind::Calls).collect();
891 assert_eq!(calls.len(), 1);
892 }
893
894 #[test]
895 fn detects_base_class_implements() {
896 let src = "class Base:\n pass\nclass Child(Base):\n pass\n";
897 let (_, _, _, _, implements, _) = parse_full(src);
898 assert!(
899 implements.iter().any(|(_, name)| name == "Base"),
900 "expected Implements edge to Base, got: {implements:?}"
901 );
902 }
903
904 #[test]
905 fn detects_type_annotation_uses() {
906 let src = "class Foo:\n pass\ndef bar(x: Foo) -> Foo:\n pass\n";
907 let (_, _, _, uses, _, _) = parse_full(src);
908 let uses_foo: Vec<_> = uses.iter().filter(|(_, n)| n == "Foo").collect();
909 assert!(
910 uses_foo.len() >= 2,
911 "expected at least 2 Uses edges to Foo, got: {uses_foo:?}"
912 );
913 }
914
915 #[test]
916 fn detects_decorator_uses() {
917 let src = "class Foo:\n @property\n def name(self):\n return self._name\n";
918 let (_, _, _, uses, _, _) = parse_full(src);
919 assert!(
920 uses.iter().any(|(_, n)| n == "property"),
921 "expected Uses edge to 'property' decorator, got: {uses:?}"
922 );
923 }
924
925 #[test]
926 fn detects_import_statement() {
927 let src = "import os\nimport sys\n\ndef main():\n pass\n";
928 let (_, _, _, _, _, imports) = parse_full(src);
929 assert!(
930 imports.iter().any(|(_, n)| n == "os"),
931 "expected import 'os', got: {imports:?}"
932 );
933 assert!(
934 imports.iter().any(|(_, n)| n == "sys"),
935 "expected import 'sys', got: {imports:?}"
936 );
937 }
938
939 #[test]
940 fn detects_from_import_statement() {
941 let src = "from os.path import join, exists\n\ndef main():\n pass\n";
942 let (_, _, _, _, _, imports) = parse_full(src);
943 assert!(
944 imports.iter().any(|(_, n)| n == "join"),
945 "expected import 'join', got: {imports:?}"
946 );
947 assert!(
948 imports.iter().any(|(_, n)| n == "exists"),
949 "expected import 'exists', got: {imports:?}"
950 );
951 }
952
953 #[test]
954 fn module_node_is_emitted() {
955 let (nodes, _) = parse("x = 1\n");
956 let modules: Vec<_> = nodes
957 .iter()
958 .filter(|n| n.kind == NodeKind::Module)
959 .collect();
960 assert_eq!(modules.len(), 1);
961 assert_eq!(modules[0].name, "test");
962 }
963
964 #[test]
965 fn async_function_flagged() {
966 let src = "async def fetch():\n pass\n";
967 let (nodes, _) = parse(src);
968 let fns: Vec<_> = nodes
969 .iter()
970 .filter(|n| n.kind == NodeKind::Function)
971 .collect();
972 assert_eq!(fns.len(), 1);
973 assert!(fns[0].metadata.is_async);
974 }
975
976 #[test]
979 fn protocol_class_becomes_interface() {
980 let src = "from typing import Protocol\n\nclass MyProto(Protocol):\n def do_it(self) -> None:\n ...\n";
981 let (nodes, _) = parse(src);
982 let ifaces: Vec<_> = nodes
983 .iter()
984 .filter(|n| n.kind == NodeKind::Interface)
985 .collect();
986 assert_eq!(ifaces.len(), 1, "expected 1 Interface, got: {ifaces:?}");
987 assert_eq!(ifaces[0].name, "MyProto");
988 assert!(
989 ifaces[0].metadata.is_abstract,
990 "Protocol should be is_abstract"
991 );
992 }
993
994 #[test]
995 fn non_protocol_class_is_struct() {
996 let src = "class Plain:\n def work(self):\n pass\n";
997 let (nodes, _) = parse(src);
998 let structs: Vec<_> = nodes
999 .iter()
1000 .filter(|n| n.kind == NodeKind::Struct)
1001 .collect();
1002 assert_eq!(structs.len(), 1);
1003 assert_eq!(structs[0].name, "Plain");
1004 assert!(!structs[0].metadata.is_abstract);
1005 }
1006
1007 #[test]
1010 fn property_decorator_yields_property_kind() {
1011 let src =
1012 "class Foo:\n @property\n def bar(self) -> str:\n return self._bar\n";
1013 let (nodes, _) = parse(src);
1014 let props: Vec<_> = nodes
1015 .iter()
1016 .filter(|n| n.kind == NodeKind::Property)
1017 .collect();
1018 assert_eq!(props.len(), 1, "expected 1 Property node, got: {props:?}");
1019 assert_eq!(props[0].name, "bar");
1020 assert!(props[0].metadata.is_property);
1021 }
1022
1023 #[test]
1024 fn staticmethod_decorator_sets_is_static() {
1025 let src = "class Foo:\n @staticmethod\n def create(x: int) -> 'Foo':\n return Foo()\n";
1026 let (nodes, _) = parse(src);
1027 let methods: Vec<_> = nodes
1028 .iter()
1029 .filter(|n| n.kind == NodeKind::Method)
1030 .collect();
1031 assert_eq!(methods.len(), 1);
1032 assert_eq!(methods[0].name, "create");
1033 assert!(
1034 methods[0].metadata.is_static,
1035 "staticmethod should set is_static"
1036 );
1037 }
1038
1039 #[test]
1040 fn classmethod_decorator_sets_is_static() {
1041 let src = "class Foo:\n @classmethod\n def from_str(cls, s: str) -> 'Foo':\n return cls()\n";
1042 let (nodes, _) = parse(src);
1043 let methods: Vec<_> = nodes
1044 .iter()
1045 .filter(|n| n.kind == NodeKind::Method)
1046 .collect();
1047 assert_eq!(methods.len(), 1);
1048 assert_eq!(methods[0].name, "from_str");
1049 assert!(
1050 methods[0].metadata.is_static,
1051 "classmethod should set is_static"
1052 );
1053 }
1054
1055 #[test]
1056 fn dataclass_decorator_class_is_struct() {
1057 let src = "from dataclasses import dataclass\n\n@dataclass\nclass Point:\n x: float\n y: float\n";
1058 let (nodes, _) = parse(src);
1059 let structs: Vec<_> = nodes
1060 .iter()
1061 .filter(|n| n.kind == NodeKind::Struct)
1062 .collect();
1063 assert_eq!(structs.len(), 1, "expected Struct for @dataclass Point");
1064 assert_eq!(structs[0].name, "Point");
1065 }
1066
1067 #[test]
1070 fn generator_function_sets_is_generator() {
1071 let src = "def numbers():\n yield 1\n yield 2\n";
1072 let (nodes, _) = parse(src);
1073 let fns: Vec<_> = nodes
1074 .iter()
1075 .filter(|n| n.kind == NodeKind::Function)
1076 .collect();
1077 assert_eq!(fns.len(), 1);
1078 assert!(
1079 fns[0].metadata.is_generator,
1080 "yield fn should be is_generator"
1081 );
1082 assert!(!fns[0].metadata.is_async);
1083 }
1084
1085 #[test]
1086 fn async_generator_is_both_async_and_generator() {
1087 let src = "async def stream():\n yield 1\n yield 2\n";
1088 let (nodes, _) = parse(src);
1089 let fns: Vec<_> = nodes
1090 .iter()
1091 .filter(|n| n.kind == NodeKind::Function)
1092 .collect();
1093 assert_eq!(fns.len(), 1);
1094 assert!(
1095 fns[0].metadata.is_async,
1096 "async generator should be is_async"
1097 );
1098 assert!(
1099 fns[0].metadata.is_generator,
1100 "async generator should be is_generator"
1101 );
1102 }
1103
1104 #[test]
1105 fn nested_yield_does_not_pollute_outer_function() {
1106 let src = "def outer():\n def inner():\n yield 1\n return inner()\n";
1108 let (nodes, _) = parse(src);
1109 let fns: Vec<_> = nodes
1110 .iter()
1111 .filter(|n| n.kind == NodeKind::Function)
1112 .collect();
1113 let outer = fns
1114 .iter()
1115 .find(|n| n.name == "outer")
1116 .expect("outer not found");
1117 assert!(
1118 !outer.metadata.is_generator,
1119 "outer should NOT be generator — yield is in nested fn"
1120 );
1121 }
1122
1123 #[test]
1126 fn module_level_bindings_detected() {
1127 let src = "MAX_SIZE = 100\nDEFAULT_NAME = 'anon'\ndefault_config = {}\n__version__ = '1.0'\nT = 1\n\ndef f():\n local = 1\n return local\n";
1131 let (nodes, _) = parse(src);
1132 let names: Vec<&str> = nodes
1133 .iter()
1134 .filter(|n| n.kind == NodeKind::Constant)
1135 .map(|n| n.name.as_str())
1136 .collect();
1137 for expected in [
1138 "MAX_SIZE",
1139 "DEFAULT_NAME",
1140 "default_config",
1141 "__version__",
1142 "T",
1143 ] {
1144 assert!(
1145 names.contains(&expected),
1146 "expected module-level binding `{expected}` as Constant; got {names:?}"
1147 );
1148 }
1149 assert!(
1150 !names.contains(&"local"),
1151 "function-local assignment must not be captured as a module Constant"
1152 );
1153 }
1154
1155 #[test]
1158 fn nested_class_emits_contains_edge_from_parent() {
1159 let src = "class Outer:\n class Inner:\n pass\n";
1160 let (nodes, edges) = parse(src);
1161 let outer = nodes
1162 .iter()
1163 .find(|n| n.name == "Outer")
1164 .expect("Outer not found");
1165 let inner = nodes
1166 .iter()
1167 .find(|n| n.name == "Inner")
1168 .expect("Inner not found");
1169 let contains: Vec<_> = edges
1170 .iter()
1171 .filter(|e| e.kind == EdgeKind::Contains)
1172 .collect();
1173 assert!(
1174 contains
1175 .iter()
1176 .any(|e| e.src == outer.id && e.dst == inner.id),
1177 "expected Contains Outer → Inner"
1178 );
1179 }
1180
1181 #[test]
1184 fn multiple_type_annotations_produce_uses_entries() {
1185 let src =
1186 "class Req:\n pass\nclass Resp:\n pass\ndef handler(r: Req) -> Resp:\n pass\n";
1187 let (_, _, _, uses, _, _) = parse_full(src);
1188 let uses_req: Vec<_> = uses.iter().filter(|(_, n)| n == "Req").collect();
1189 let uses_resp: Vec<_> = uses.iter().filter(|(_, n)| n == "Resp").collect();
1190 assert!(!uses_req.is_empty(), "expected Uses edge to Req");
1191 assert!(!uses_resp.is_empty(), "expected Uses edge to Resp");
1192 }
1193
1194 #[test]
1197 fn private_method_has_private_visibility() {
1198 use gitcortex_core::schema::Visibility;
1199 let src = "class Foo:\n def _internal(self):\n pass\n def public(self):\n pass\n";
1200 let (nodes, _) = parse(src);
1201 let internal = nodes
1202 .iter()
1203 .find(|n| n.name == "_internal")
1204 .expect("_internal not found");
1205 let public = nodes
1206 .iter()
1207 .find(|n| n.name == "public")
1208 .expect("public not found");
1209 assert_eq!(internal.metadata.visibility, Visibility::Private);
1210 assert_eq!(public.metadata.visibility, Visibility::Pub);
1211 }
1212
1213 #[test]
1216 fn calls_edge_between_two_functions() {
1217 let src = "def helper():\n pass\n\ndef main():\n helper()\n";
1218 let (_, edges) = parse(src);
1219 let calls: Vec<_> = edges.iter().filter(|e| e.kind == EdgeKind::Calls).collect();
1220 assert_eq!(calls.len(), 1, "expected exactly 1 Calls edge");
1221 }
1222
1223 #[test]
1224 fn method_call_via_self_creates_calls_edge() {
1225 let src = "class Svc:\n def run(self):\n self.process()\n def process(self):\n pass\n";
1226 let (nodes, edges) = parse(src);
1227 let run = nodes
1228 .iter()
1229 .find(|n| n.name == "run")
1230 .expect("run not found");
1231 let process = nodes
1232 .iter()
1233 .find(|n| n.name == "process")
1234 .expect("process not found");
1235 let calls: Vec<_> = edges.iter().filter(|e| e.kind == EdgeKind::Calls).collect();
1236 assert!(
1238 calls.iter().any(|e| e.src == run.id && e.dst == process.id),
1239 "expected Calls edge run → process, got: {calls:?}"
1240 );
1241 }
1242
1243 #[test]
1246 fn aliased_import_uses_alias_name() {
1247 let src = "import numpy as np\nimport pandas as pd\n";
1248 let (_, _, _, _, _, imports) = parse_full(src);
1249 let names: Vec<&str> = imports.iter().map(|(_, n)| n.as_str()).collect();
1251 assert!(
1252 names.contains(&"numpy"),
1253 "expected import 'numpy', got: {names:?}"
1254 );
1255 assert!(
1256 names.contains(&"pandas"),
1257 "expected import 'pandas', got: {names:?}"
1258 );
1259 }
1260
1261 #[test]
1262 fn dotted_import_records_leaf_module() {
1263 let src = "import os.path\n";
1264 let (_, _, _, _, _, imports) = parse_full(src);
1265 let names: Vec<&str> = imports.iter().map(|(_, n)| n.as_str()).collect();
1266 assert!(
1267 names.contains(&"path"),
1268 "expected leaf 'path' from 'import os.path', got: {names:?}"
1269 );
1270 }
1271}