1use tree_sitter::{Language, Parser, Tree};
2
3use crate::core::code_graph::CodeGraph;
4use crate::core::moniker::Moniker;
5use crate::core::shape::Shape;
6
7use crate::lang::{ExtractionContext, KindSpec, LangExtractor, ParsedDocument};
8
9pub mod build;
10mod canonicalize;
11mod kinds;
12mod sdk_pipeline;
13
14#[derive(Clone, Debug, Default)]
15pub struct Presets {}
16
17pub fn parse(source: &str) -> Tree {
18 let mut parser = Parser::new();
19 let language: Language = tree_sitter_python::LANGUAGE.into();
20 parser.set_language(&language).unwrap_or_else(|err| {
21 panic!("failed to load tree-sitter Python grammar: {err}");
22 });
23 parser.parse(source, None).unwrap_or_else(|| {
24 panic!("tree-sitter parse returned None on a non-cancelled call");
25 })
26}
27
28pub fn extract(
29 uri: &str,
30 source: &str,
31 anchor: &Moniker,
32 deep: bool,
33 presets: &Presets,
34) -> CodeGraph {
35 <Lang as LangExtractor>::extract(uri, source, anchor, deep, presets)
36}
37
38pub struct Lang;
39
40const DEF_KINDS: &[&str] = &[
41 "class",
42 "type",
43 "function",
44 "method",
45 "async_function",
46 "path",
47];
48
49const DEF_KIND_SPECS: &[KindSpec] = &[
50 KindSpec::new("class", Shape::Type, 20, "class"),
51 KindSpec::new("type", Shape::Type, 21, "type"),
52 KindSpec::new("function", Shape::Callable, 40, "function"),
53 KindSpec::new("async_function", Shape::Callable, 41, "async_function"),
54 KindSpec::new("method", Shape::Callable, 42, "method"),
55 KindSpec::new("path", Shape::Value, 63, "path"),
56];
57
58impl crate::lang::LangExtractor for Lang {
59 type Presets = Presets;
60 const LANG_TAG: &'static str = "python";
61 const ALLOWED_KINDS: &'static [&'static str] = DEF_KINDS;
62 const KIND_SPECS: &'static [KindSpec] = DEF_KIND_SPECS;
63 const ALLOWED_VISIBILITIES: &'static [&'static str] = &["public", "private", "module"];
64
65 fn parse(_uri: &str, source: &str) -> ParsedDocument {
66 ParsedDocument::new(parse(source))
67 }
68
69 fn file_root(uri: &str, anchor: &Moniker) -> Option<Moniker> {
70 Some(canonicalize::compute_module_moniker(anchor, uri))
71 }
72
73 fn extract_parsed(
74 context: ExtractionContext<'_, Self::Presets>,
75 document: &ParsedDocument,
76 ) -> CodeGraph {
77 sdk_pipeline::extract(
78 context.uri,
79 context.source,
80 document,
81 context.anchor,
82 context.deep,
83 context.presets,
84 )
85 }
86}
87
88#[cfg(test)]
89mod tests {
90 use super::*;
91 use crate::core::moniker::MonikerBuilder;
92 use crate::lang::assert_conformance;
93
94 fn make_anchor() -> Moniker {
95 MonikerBuilder::new().project(b"app").build()
96 }
97
98 fn extract_default(uri: &str, source: &str, anchor: &Moniker, deep: bool) -> CodeGraph {
99 let g = extract(uri, source, anchor, deep, &Presets::default());
100 assert_conformance::<super::Lang>(&g, anchor);
101 g
102 }
103
104 #[test]
105 fn parse_empty_returns_module() {
106 let tree = parse("");
107 assert_eq!(tree.root_node().kind(), "module");
108 }
109
110 #[test]
111 fn extract_module_uses_path_segments() {
112 let g = extract_default("acme/util/text.py", "", &make_anchor(), false);
113 let expected = MonikerBuilder::new()
114 .project(b"app")
115 .segment(b"lang", b"python")
116 .segment(b"package", b"acme")
117 .segment(b"package", b"util")
118 .segment(b"module", b"text")
119 .build();
120 assert_eq!(g.root(), &expected);
121 }
122
123 #[test]
124 fn extract_module_root_is_filename_only() {
125 let g = extract_default("foo.py", "", &make_anchor(), false);
126 let expected = MonikerBuilder::new()
127 .project(b"app")
128 .segment(b"lang", b"python")
129 .segment(b"module", b"foo")
130 .build();
131 assert_eq!(g.root(), &expected);
132 }
133
134 #[test]
135 fn extract_function_with_typed_params_emits_full_signature() {
136 let src = "def make(x: int, y: str) -> int:\n return x\n";
137 let g = extract_default("m.py", src, &make_anchor(), false);
138 let f = g
139 .defs()
140 .find(|d| d.kind == b"function")
141 .expect("function def");
142 let last = f.moniker.as_view().segments().last().unwrap();
143 assert_eq!(last.kind, b"function");
144 assert_eq!(last.name, b"make(x:int,y:str)");
145 assert_eq!(f.signature, b"x:int,y:str".to_vec());
146 }
147
148 #[test]
149 fn extract_function_with_untyped_params_uses_name_only_slots() {
150 let src = "def f(a, b=1):\n return a\n";
151 let g = extract_default("m.py", src, &make_anchor(), false);
152 let f = g
153 .defs()
154 .find(|d| d.kind == b"function")
155 .expect("function def");
156 let last = f.moniker.as_view().segments().last().unwrap();
157 assert_eq!(last.name, b"f(a,b)");
158 assert_eq!(f.signature, b"a,b".to_vec());
159 }
160
161 #[test]
162 fn extract_classmethod_excludes_cls_from_signature() {
163 let src = "class Foo:\n @classmethod\n def make(cls, x: int) -> 'Foo':\n return cls()\n";
164 let g = extract_default("foo.py", src, &make_anchor(), false);
165 let m = g.defs().find(|d| d.kind == b"method").expect("method def");
166 assert_eq!(
167 m.moniker.as_view().segments().last().unwrap().name,
168 b"make(x:int)"
169 );
170 }
171
172 #[test]
173 fn extract_double_underscore_visibility_is_private() {
174 let src = "class Foo:\n def __secret(self):\n pass\n";
175 let g = extract_default("foo.py", src, &make_anchor(), false);
176 let m = g.defs().find(|d| d.kind == b"method").expect("method def");
177 assert_eq!(m.visibility, b"private".to_vec());
178 }
179
180 #[test]
181 fn extract_single_underscore_visibility_is_module() {
182 let src = "def _internal():\n pass\n";
183 let g = extract_default("m.py", src, &make_anchor(), false);
184 let f = g
185 .defs()
186 .find(|d| d.kind == b"function")
187 .expect("function def");
188 assert_eq!(f.visibility, b"module".to_vec());
189 }
190
191 #[test]
192 fn extract_import_module_emits_imports_module() {
193 let src = "import os\nimport acme.util as u\n";
194 let g = extract_default("m.py", src, &make_anchor(), false);
195 let kinds: Vec<&[u8]> = g.refs().map(|r| r.kind.as_ref()).collect();
196 assert_eq!(kinds.iter().filter(|k| **k == b"imports_module").count(), 2);
197 }
198
199 #[test]
200 fn extract_stdlib_import_marks_external() {
201 let g = extract_default("m.py", "import json\n", &make_anchor(), false);
202 let r = g
203 .refs()
204 .find(|r| r.kind == b"imports_module")
205 .expect("imports_module");
206 assert_eq!(r.confidence, b"external".to_vec());
207 }
208
209 #[test]
210 fn extract_project_import_marks_imported() {
211 let g = extract_default("m.py", "import acme.util\n", &make_anchor(), false);
212 let r = g
213 .refs()
214 .find(|r| r.kind == b"imports_module")
215 .expect("imports_module");
216 assert_eq!(r.confidence, b"imported".to_vec());
217 }
218
219 #[test]
220 fn extract_from_import_emits_one_imports_symbol_per_name() {
221 let src = "from acme.util import a, b as c\n";
222 let g = extract_default("m.py", src, &make_anchor(), false);
223 let names: Vec<&[u8]> = g
224 .refs()
225 .filter(|r| r.kind == b"imports_symbol")
226 .map(|r| r.target.as_view().segments().last().unwrap().name)
227 .collect();
228 assert_eq!(names, vec![&b"a"[..], &b"b"[..]]);
229 let segs: Vec<_> = g
230 .refs()
231 .find(|r| r.kind == b"imports_symbol")
232 .unwrap()
233 .target
234 .as_view()
235 .segments()
236 .collect();
237 let kinds: Vec<&[u8]> = segs.iter().map(|s| s.kind).collect();
238 assert_eq!(
239 kinds,
240 vec![&b"lang"[..], &b"package"[..], &b"module"[..], &b"path"[..]]
241 );
242 let aliased = g
243 .refs()
244 .find(|r| r.kind == b"imports_symbol" && r.alias == b"c")
245 .expect("aliased import");
246 assert_eq!(aliased.alias, b"c".to_vec());
247 }
248
249 #[test]
250 fn extract_future_annotations_does_not_emit_a_runtime_read() {
251 let g = extract_default(
252 "m.py",
253 "from __future__ import annotations\n",
254 &make_anchor(),
255 false,
256 );
257 assert!(!g.refs().any(|reference| {
258 reference.kind == b"reads"
259 && reference.target.as_view().segments().last().unwrap().name == b"annotations"
260 }));
261 }
262
263 #[test]
264 fn extract_runtime_module_globals_are_external() {
265 let g = extract_default(
266 "m.py",
267 "if __name__ == '__main__':\n pass\n",
268 &make_anchor(),
269 false,
270 );
271 let reference = g
272 .refs()
273 .find(|reference| reference.kind == b"reads")
274 .expect("reads __name__");
275 assert_eq!(reference.confidence, b"external".to_vec());
276 assert_eq!(
277 reference.target.as_view().segments().next().unwrap().name,
278 b"python"
279 );
280 assert_eq!(
281 reference.target.as_view().segments().nth(1).unwrap().name,
282 b"runtime"
283 );
284 }
285
286 #[test]
287 fn extract_wildcard_import_preserves_the_star_binding() {
288 let g = extract_default(
289 "acme/facade.py",
290 "from .implementation import *\n",
291 &make_anchor(),
292 false,
293 );
294 let reference = g
295 .refs()
296 .find(|reference| reference.kind == b"imports_module")
297 .expect("wildcard module import");
298
299 assert_eq!(reference.alias, b"*".to_vec());
300 }
301
302 #[test]
303 fn extract_static_all_emits_named_export_bindings() {
304 let src = "__all__ = [\n \"PublicClient\",\n # Kept public for compatibility.\n 'helper',\n]\nclass PublicClient:\n pass\ndef helper():\n pass\n";
305 let g = extract_default("acme/implementation.py", src, &make_anchor(), false);
306 let exports = g
307 .refs()
308 .filter(|reference| reference.kind == b"reexports" && !reference.alias.is_empty())
309 .map(|reference| reference.alias.clone())
310 .collect::<Vec<_>>();
311
312 assert_eq!(exports, vec![b"PublicClient".to_vec(), b"helper".to_vec()]);
313 }
314
315 #[test]
316 fn extract_all_assignments_preserve_empty_dynamic_and_extend_state() {
317 let src = "__all__ = []\n__all__ += ['Later']\n__all__ = build_exports()\n";
318 let g = extract_default("acme/facade.py", src, &make_anchor(), false);
319 let directives = g
320 .refs()
321 .filter(|reference| reference.kind == b"reexports" && reference.alias.is_empty())
322 .map(|reference| reference.receiver_hint.clone())
323 .collect::<Vec<_>>();
324
325 assert_eq!(
326 directives,
327 vec![
328 b"python_all_replace".to_vec(),
329 b"python_all_extend".to_vec(),
330 b"python_all_dynamic".to_vec(),
331 ]
332 );
333 }
334
335 #[test]
336 fn extract_conditional_import_marks_runtime_binding() {
337 let src = "if enabled:\n from acme.client import Client\n";
338 let g = extract_default("acme/facade.py", src, &make_anchor(), false);
339 let reference = g
340 .refs()
341 .find(|reference| reference.kind == b"imports_symbol")
342 .expect("conditional import");
343
344 assert_eq!(
345 reference.receiver_hint,
346 b"python_conditional_import".to_vec()
347 );
348 }
349
350 #[test]
351 fn extract_calls_through_conditional_imports_as_runtime_bindings() {
352 let src = "def build(enabled):\n if enabled:\n from acme.client import Client\n return Client()\n";
353 let g = extract_default("acme/facade.py", src, &make_anchor(), false);
354 let reference = g
355 .refs()
356 .find(|reference| reference.kind == b"calls")
357 .expect("call through conditional import");
358
359 assert_eq!(
360 reference.receiver_hint,
361 b"python_conditional_import".to_vec()
362 );
363 }
364
365 #[test]
366 fn extract_function_imports_do_not_leak_into_sibling_scopes() {
367 let src = "def configure():\n from acme.client import Client\n return Client()\n\ndef build():\n return Client()\n";
368 let g = extract_default("acme/facade.py", src, &make_anchor(), false);
369 let calls = g
370 .refs()
371 .filter(|reference| reference.kind == b"calls")
372 .map(|reference| {
373 let source = g
374 .def_at(reference.source)
375 .moniker
376 .as_view()
377 .segments()
378 .last()
379 .expect("call source")
380 .name
381 .to_vec();
382 (source, reference.confidence.to_vec())
383 })
384 .collect::<Vec<_>>();
385
386 assert!(
387 calls.contains(&(b"configure()".to_vec(), b"imported".to_vec())),
388 "{calls:?}"
389 );
390 assert!(
391 calls.contains(&(b"build()".to_vec(), b"name_match".to_vec())),
392 "{calls:?}"
393 );
394 }
395
396 #[test]
397 fn extract_local_import_shadows_conditional_module_binding() {
398 let src = "try:\n from acme.a import Client\nexcept ImportError:\n from acme.b import Client\n\ndef build():\n from acme.c import Client\n return Client()\n";
399 let g = extract_default("acme/facade.py", src, &make_anchor(), false);
400 let call = g
401 .refs()
402 .find(|reference| reference.kind == b"calls")
403 .expect("locally shadowed call");
404
405 assert_eq!(call.receiver_hint, b"");
406 assert_eq!(call.confidence, b"imported");
407 assert!(
408 call.target
409 .as_view()
410 .segments()
411 .any(|segment| { segment.kind == b"module" && segment.name == b"c" })
412 );
413 }
414
415 #[test]
416 fn extract_conditional_all_as_dynamic() {
417 let src = "if enabled:\n __all__ = ['Client']\nelse:\n __all__ = ['Fallback']\n";
418 let g = extract_default("acme/facade.py", src, &make_anchor(), false);
419 let directives = g
420 .refs()
421 .filter(|reference| reference.kind == b"reexports")
422 .map(|reference| (reference.alias.to_vec(), reference.receiver_hint.to_vec()))
423 .collect::<Vec<_>>();
424
425 assert_eq!(
426 directives,
427 vec![
428 (Vec::new(), b"python_all_dynamic".to_vec()),
429 (Vec::new(), b"python_all_dynamic".to_vec()),
430 ]
431 );
432 }
433
434 #[test]
435 fn extract_relative_import_resolves_against_importer() {
436 let src = "from .util import helper\n";
437 let g = extract_default("acme/m.py", src, &make_anchor(), false);
438 let r = g
439 .refs()
440 .find(|r| r.kind == b"imports_symbol")
441 .expect("imports_symbol");
442 let segs: Vec<_> = r.target.as_view().segments().collect();
443 let kinds: Vec<&[u8]> = segs.iter().map(|s| s.kind).collect();
444 let names: Vec<&[u8]> = segs.iter().map(|s| s.name).collect();
445 assert_eq!(
446 kinds,
447 vec![&b"lang"[..], &b"package"[..], &b"module"[..], &b"path"[..]]
448 );
449 assert_eq!(
450 names,
451 vec![&b"python"[..], &b"acme"[..], &b"util"[..], &b"helper"[..]]
452 );
453 }
454
455 #[test]
456 fn extract_relative_import_underflow_falls_back_to_external_pkg() {
457 let src = "from ...foo import bar\n";
458 let g = extract_default("m.py", src, &make_anchor(), false);
459 let r = g
460 .refs()
461 .find(|r| r.kind == b"imports_symbol")
462 .expect("imports_symbol");
463 let segs: Vec<_> = r.target.as_view().segments().collect();
464 assert_eq!(segs[0].kind, b"external_pkg");
465 assert_eq!(segs[0].name, b"...");
466 }
467
468 #[test]
469 fn extract_decorator_emits_annotates() {
470 let src = "import functools\n@functools.wraps(fn)\ndef g():\n pass\n";
471 let g = extract_default("m.py", src, &make_anchor(), false);
472 let ann = g
473 .refs()
474 .find(|r| r.kind == b"annotates")
475 .expect("annotates ref");
476 assert_eq!(
477 ann.target.as_view().segments().last().unwrap().name,
478 b"wraps"
479 );
480 }
481
482 #[test]
483 fn extract_param_read_marks_confidence_local() {
484 let src = "def f(x):\n return x\n";
485 let g = extract_default("m.py", src, &make_anchor(), true);
486 let r = g
487 .refs()
488 .find(|r| {
489 r.kind == b"reads" && r.target.as_view().segments().last().unwrap().name == b"x"
490 })
491 .expect("reads x");
492 assert_eq!(r.confidence, b"local".to_vec());
493 }
494
495 #[test]
496 fn extract_deep_emits_param_def() {
497 let src = "def f(x, y: int):\n return x\n";
498 let g = extract_default("m.py", src, &make_anchor(), true);
499 let params: Vec<&[u8]> = g
500 .defs()
501 .filter(|d| d.kind == b"param")
502 .map(|d| d.moniker.as_view().segments().last().unwrap().name)
503 .collect();
504 assert!(params.contains(&&b"x"[..]));
505 assert!(params.contains(&&b"y"[..]));
506 }
507
508 #[test]
509 fn extract_async_function_tracks_local_assignments() {
510 let src = "async def f():\n x = 1\n return x\n";
511 let g = extract_default("m.py", src, &make_anchor(), true);
512 let r = g
513 .refs()
514 .find(|r| {
515 r.kind == b"reads" && r.target.as_view().segments().last().unwrap().name == b"x"
516 })
517 .expect("reads x");
518 assert_eq!(r.confidence, b"local".to_vec());
519 }
520
521 #[test]
522 fn extract_imported_identifier_read_targets_import_moniker() {
523 let src = "import asyncio\nasync def f():\n return asyncio\n";
524 let g = extract_default("m.py", src, &make_anchor(), false);
525 let r = g
526 .refs()
527 .find(|r| r.kind == b"reads" && r.confidence == b"external")
528 .expect("external read");
529 let segs: Vec<_> = r.target.as_view().segments().collect();
530 assert_eq!(segs[0].kind, b"sdk");
531 assert_eq!(segs[0].name, b"python");
532 assert_eq!(segs[1].name, b"asyncio");
533 }
534
535 #[test]
536 fn extract_unknown_identifier_read_marks_incomplete_resolution() {
537 let src = "def f():\n return missing_name\n";
538 let g = extract_default("m.py", src, &make_anchor(), false);
539 let r = g
540 .refs()
541 .find(|r| {
542 r.kind == b"reads"
543 && r.target.as_view().segments().last().unwrap().name == b"missing_name"
544 })
545 .expect("reads missing_name");
546 assert_eq!(r.confidence, b"unresolved".to_vec());
547 }
548
549 #[test]
550 fn extract_same_module_callable_read_keeps_exact_resolution() {
551 let src = "def f():\n return callback\n\ndef callback():\n pass\n";
552 let g = extract_default("m.py", src, &make_anchor(), false);
553 let r = g
554 .refs()
555 .find(|r| {
556 r.kind == b"reads"
557 && r.target.as_view().segments().last().unwrap().name == b"callback()"
558 })
559 .expect("reads callback");
560 assert_eq!(r.confidence, b"resolved".to_vec());
561 }
562
563 #[test]
564 fn extract_module_callable_read_respects_local_shadowing() {
565 let src = "def callback():\n pass\n\ndef f(callback):\n return callback\n";
566 let g = extract_default("m.py", src, &make_anchor(), true);
567 let r = g
568 .refs()
569 .find(|r| {
570 r.kind == b"reads"
571 && r.target.as_view().segments().last().unwrap().name == b"callback"
572 })
573 .expect("reads local callback");
574 assert_eq!(r.confidence, b"local".to_vec());
575 assert_eq!(r.target.as_view().segments().last().unwrap().kind, b"local");
576 }
577
578 #[test]
579 fn extract_unknown_receiver_method_call_marks_incomplete_resolution() {
580 let src = "def f(value):\n return value.normalize()\n";
581 let g = extract_default("m.py", src, &make_anchor(), false);
582 let r = g
583 .refs()
584 .find(|r| {
585 r.kind == b"method_call"
586 && r.target.as_view().segments().last().unwrap().name == b"normalize"
587 })
588 .expect("method_call normalize");
589 assert_eq!(r.confidence, b"unresolved".to_vec());
590 }
591
592 #[test]
593 fn extract_self_member_call_resolves_via_typed_constructor_param() {
594 let src = "class Store:\n def reserve(self) -> None:\n pass\n\nclass Worker:\n def __init__(self, store: Store) -> None:\n self._store = store\n\n def run(self) -> None:\n self._store.reserve()\n";
595 let g = extract_default("m.py", src, &make_anchor(), false);
596 let r = g
597 .refs()
598 .find(|r| {
599 r.kind == b"method_call"
600 && r.target.as_view().segments().last().unwrap().name == b"reserve()"
601 })
602 .expect("method_call Store.reserve");
603 assert_eq!(r.confidence, b"resolved".to_vec());
604 let parent = r.target.parent().expect("method parent");
605 assert_eq!(parent.as_view().segments().last().unwrap().name, b"Store");
606 }
607
608 #[test]
609 fn extract_self_member_call_uses_constructor_type_independent_of_method_order() {
610 let src = "class Store:\n def reserve(self) -> None:\n pass\n\nclass Other:\n def reserve(self) -> None:\n pass\n\nclass Worker:\n def before_init(self) -> None:\n self._store.reserve()\n\n def __init__(self, store: Store) -> None:\n self._store = store\n\n def retarget(self, other: Other) -> None:\n self._store = other\n\n def after_retarget(self) -> None:\n self._store.reserve()\n";
611 let g = extract_default("m.py", src, &make_anchor(), false);
612 let targets: Vec<Moniker> = g
613 .refs()
614 .filter(|r| {
615 r.kind == b"method_call"
616 && r.target.as_view().segments().last().unwrap().name == b"reserve()"
617 })
618 .map(|r| r.target.clone())
619 .collect();
620 assert_eq!(targets.len(), 2);
621 for target in targets {
622 let parent = target.parent().expect("method parent");
623 assert_eq!(parent.as_view().segments().last().unwrap().name, b"Store");
624 }
625 }
626
627 #[test]
628 fn extract_self_callable_attr_call_targets_attribute_type_alias() {
629 let src = "from collections.abc import Callable\nfrom typing import TypeAlias\n\nCallback: TypeAlias = Callable[[], None]\n\nclass Worker:\n def __init__(self, cb: Callback) -> None:\n self._cb = cb\n\n def run(self) -> None:\n self._cb()\n";
630 let g = extract_default("m.py", src, &make_anchor(), false);
631 let r = g
632 .refs()
633 .find(|r| {
634 r.kind == b"calls"
635 && r.target.as_view().segments().last().unwrap().name == b"Callback"
636 })
637 .expect("calls Callback type alias");
638 assert_eq!(r.confidence, b"resolved".to_vec());
639 assert!(!g.refs().any(|r| {
640 r.kind == b"method_call" && r.target.as_view().segments().last().unwrap().name == b"_cb"
641 }));
642 }
643
644 #[test]
645 fn extract_type_alias_emits_type_def_and_rhs_uses_type() {
646 let src = "from typing import TypeAlias\n\nclass User:\n pass\n\nUserMap: TypeAlias = dict[str, User]\ntype UserResult = User | None\n";
647 let g = extract_default("m.py", src, &make_anchor(), false);
648 let type_names: Vec<&[u8]> = g
649 .defs()
650 .filter(|d| d.kind == b"type")
651 .map(|d| d.moniker.as_view().segments().last().unwrap().name)
652 .collect();
653 assert!(type_names.contains(&&b"UserMap"[..]));
654 assert!(type_names.contains(&&b"UserResult"[..]));
655 assert!(g.refs().any(|r| {
656 r.kind == b"uses_type"
657 && r.target.as_view().segments().last().unwrap().name == b"User"
658 && r.confidence == b"resolved"
659 }));
660 assert!(!g.refs().any(|r| {
661 r.kind == b"uses_type"
662 && matches!(
663 r.target.as_view().segments().last().unwrap().name,
664 b"dict" | b"str" | b"None"
665 )
666 }));
667 }
668
669 #[test]
670 fn extract_local_class_call_emits_instantiates() {
671 let src = "class User:\n pass\n\ndef make() -> User:\n return User()\n";
672 let g = extract_default("m.py", src, &make_anchor(), false);
673 let r = g
674 .refs()
675 .find(|r| r.kind == b"instantiates")
676 .expect("instantiates User");
677 assert_eq!(r.confidence, b"resolved".to_vec());
678 assert_eq!(r.target.as_view().segments().last().unwrap().name, b"User");
679 }
680
681 #[test]
682 fn extract_callable_return_annotation_emits_returns_type() {
683 let src = "class User:\n pass\n\ndef make() -> User:\n return User()\n";
684 let g = extract_default("m.py", src, &make_anchor(), false);
685 let r = g
686 .refs()
687 .find(|r| r.kind == b"returns_type")
688 .expect("returns_type User");
689 assert_eq!(r.confidence, b"resolved".to_vec());
690 assert_eq!(r.target.as_view().segments().last().unwrap().kind, b"class");
691 assert_eq!(r.target.as_view().segments().last().unwrap().name, b"User");
692 }
693
694 #[test]
695 fn extract_local_assignment_uses_annotated_factory_return_type() {
696 let src = "class User:\n def label(self) -> str:\n return 'user'\n\ndef make() -> User:\n return User()\n\ndef render():\n value = make()\n return value.label()\n";
697 let g = extract_default("m.py", src, &make_anchor(), false);
698 let r = g
699 .refs()
700 .find(|r| {
701 r.kind == b"method_call"
702 && r.target.as_view().segments().last().unwrap().name == b"label()"
703 })
704 .expect("method_call User.label");
705 assert_eq!(r.confidence, b"resolved".to_vec());
706 assert_eq!(
707 r.target
708 .parent()
709 .expect("method parent")
710 .as_view()
711 .segments()
712 .last()
713 .unwrap()
714 .name,
715 b"User"
716 );
717 }
718
719 #[test]
720 fn extract_shadowed_factory_does_not_leak_module_return_type() {
721 let src = "class User:\n def label(self) -> str:\n return 'user'\n\ndef make() -> User:\n return User()\n\ndef render(make):\n value = make()\n return value.label()\n";
722 let g = extract_default("m.py", src, &make_anchor(), false);
723 let r = g
724 .refs()
725 .find(|r| {
726 r.kind == b"method_call"
727 && r.target.as_view().segments().last().unwrap().name == b"label"
728 })
729 .expect("method_call label");
730 assert_eq!(r.confidence, b"unresolved".to_vec());
731 }
732
733 #[test]
734 fn extract_later_assignment_still_shadows_module_factory() {
735 let src = "class User:\n def label(self) -> str:\n return 'user'\n\ndef make() -> User:\n return User()\n\ndef render():\n value = make()\n make = lambda: None\n return value.label()\n";
736 let g = extract_default("m.py", src, &make_anchor(), false);
737 let r = g
738 .refs()
739 .find(|r| {
740 r.kind == b"method_call"
741 && r.target.as_view().segments().last().unwrap().name == b"label"
742 })
743 .expect("method_call label");
744 assert_eq!(r.confidence, b"unresolved".to_vec());
745 }
746
747 #[test]
748 fn extract_factory_return_type_uses_defining_scope() {
749 let src = "class Product:\n def label(self) -> str:\n return 'product'\n\ndef make() -> Product:\n return Product()\n\nclass View:\n class Product:\n def label(self) -> str:\n return 'shadow'\n\n def render(self):\n value = make()\n return value.label()\n";
750 let g = extract_default("m.py", src, &make_anchor(), false);
751 let r = g
752 .refs()
753 .find(|r| {
754 r.kind == b"method_call"
755 && g.defs().nth(r.source).is_some_and(|source| {
756 source
757 .moniker
758 .as_view()
759 .segments()
760 .last()
761 .is_some_and(|segment| segment.name == b"render()")
762 }) && r.target.as_view().segments().last().unwrap().name == b"label()"
763 })
764 .expect("method_call Product.label");
765 let owner = r.target.parent().expect("method owner");
766 assert_eq!(owner.as_view().segments().last().unwrap().name, b"Product");
767 assert_eq!(owner.parent(), Some(g.root().clone()));
768 }
769
770 #[test]
771 fn extract_bare_call_in_method_does_not_bind_sibling_method() {
772 let src = "class Product:\n def label(self):\n pass\n\nclass Factory:\n def make(self) -> Product:\n return Product()\n\n def render(self):\n value = make()\n return value.label()\n";
773 let g = extract_default("m.py", src, &make_anchor(), false);
774 let r = g
775 .refs()
776 .find(|r| {
777 r.kind == b"method_call"
778 && r.target.as_view().segments().last().unwrap().name == b"label"
779 })
780 .expect("method_call label");
781 assert_eq!(r.confidence, b"unresolved".to_vec());
782 }
783
784 #[test]
785 fn extract_builtin_annotated_receiver_method_is_external() {
786 let src = "def normalize(value: str) -> str:\n return value.strip()\n";
787 let g = extract_default("m.py", src, &make_anchor(), false);
788 let r = g
789 .refs()
790 .find(|r| {
791 r.kind == b"method_call"
792 && r.target.as_view().segments().last().unwrap().name == b"strip"
793 })
794 .expect("method_call str.strip");
795 assert_eq!(r.confidence, b"external".to_vec());
796 let segments = r.target.as_view().segments().collect::<Vec<_>>();
797 assert_eq!(segments[0].kind, b"sdk");
798 assert_eq!(segments[0].name, b"python");
799 assert_eq!(segments[1].name, b"builtins");
800 assert_eq!(segments[2].name, b"str");
801 }
802
803 #[test]
804 fn extract_builtin_generic_receiver_uses_container_type() {
805 let src = "def add(values: list[str]):\n values.append('x')\n";
806 let g = extract_default("m.py", src, &make_anchor(), false);
807 let r = g
808 .refs()
809 .find(|reference| reference.kind == b"method_call")
810 .expect("method_call list.append");
811 assert_eq!(r.confidence, b"external".to_vec());
812 assert_eq!(r.target.as_view().segments().nth(2).unwrap().name, b"list");
813 }
814
815 #[test]
816 fn extract_variadic_parameters_use_runtime_container_types() {
817 let src = "def collect(*items: str, **options: int):\n items.count('x')\n options.get('limit')\n";
818 let g = extract_default("m.py", src, &make_anchor(), false);
819 let targets = g
820 .refs()
821 .filter(|reference| reference.kind == b"method_call")
822 .map(|reference| {
823 reference
824 .target
825 .as_view()
826 .segments()
827 .nth(2)
828 .unwrap()
829 .name
830 .to_vec()
831 })
832 .collect::<Vec<_>>();
833 assert!(targets.contains(&b"tuple".to_vec()), "{targets:?}");
834 assert!(targets.contains(&b"dict".to_vec()), "{targets:?}");
835 }
836
837 #[test]
838 fn extract_with_alias_preserves_constructed_receiver_type() {
839 let src = "class Session:\n def __enter__(self) -> Session:\n return self\n\n def __exit__(self, exc_type, exc, tb):\n pass\n\n def close(self):\n pass\n\ndef run():\n with Session() as session:\n session.close()\n";
840 let g = extract_default("m.py", src, &make_anchor(), false);
841 let r = g
842 .refs()
843 .find(|reference| reference.kind == b"method_call")
844 .expect("method_call Session.close");
845 assert_eq!(r.confidence, b"resolved".to_vec());
846 assert_eq!(
847 r.target
848 .parent()
849 .unwrap()
850 .as_view()
851 .segments()
852 .last()
853 .unwrap()
854 .name,
855 b"Session"
856 );
857 }
858
859 #[test]
860 fn extract_annotated_iterable_types_the_loop_binding() {
861 let src = "class Item:\n def label(self):\n pass\n\ndef render(items: list[Item]):\n for item in items:\n item.label()\n";
862 let g = extract_default("m.py", src, &make_anchor(), false);
863 let r = g
864 .refs()
865 .find(|reference| reference.kind == b"method_call")
866 .expect("method_call Item.label");
867 assert_eq!(r.confidence, b"resolved".to_vec());
868 assert_eq!(
869 r.target
870 .parent()
871 .unwrap()
872 .as_view()
873 .segments()
874 .last()
875 .unwrap()
876 .name,
877 b"Item"
878 );
879 }
880
881 #[test]
882 fn extract_heterogeneous_tuple_annotation_keeps_all_loop_element_types() {
883 let src = "class Alpha:\n pass\n\nclass Beta:\n pass\n\ndef render(values: tuple[Alpha, Beta]):\n for value in values:\n pass\n";
884 let g = extract_default("m.py", src, &make_anchor(), false);
885 let mut targets = g
886 .refs()
887 .filter(|reference| reference.kind == b"typed_as" && reference.alias == b"value")
888 .map(|reference| {
889 reference
890 .target
891 .as_view()
892 .segments()
893 .last()
894 .unwrap()
895 .name
896 .to_vec()
897 })
898 .collect::<Vec<_>>();
899 targets.sort();
900 assert_eq!(targets, vec![b"Alpha".to_vec(), b"Beta".to_vec()]);
901 }
902
903 #[test]
904 fn extract_except_alias_uses_the_exception_type() {
905 let src = "class Problem(Exception):\n def explain(self):\n pass\n\ndef run():\n try:\n pass\n except Problem as error:\n error.explain()\n";
906 let g = extract_default("m.py", src, &make_anchor(), false);
907 let r = g
908 .refs()
909 .find(|reference| reference.kind == b"method_call")
910 .expect("method_call Problem.explain");
911 assert_eq!(r.confidence, b"resolved".to_vec());
912 assert_eq!(
913 r.target
914 .parent()
915 .unwrap()
916 .as_view()
917 .segments()
918 .last()
919 .unwrap()
920 .name,
921 b"Problem"
922 );
923 }
924
925 #[test]
926 fn extract_exception_tuple_keeps_all_alias_types() {
927 let src = "class AlphaError(Exception):\n pass\n\nclass BetaError(Exception):\n pass\n\ndef run():\n try:\n pass\n except (AlphaError, BetaError) as error:\n pass\n";
928 let g = extract_default("m.py", src, &make_anchor(), false);
929 let mut targets = g
930 .refs()
931 .filter(|reference| reference.kind == b"typed_as" && reference.alias == b"error")
932 .map(|reference| {
933 reference
934 .target
935 .as_view()
936 .segments()
937 .last()
938 .unwrap()
939 .name
940 .to_vec()
941 })
942 .collect::<Vec<_>>();
943 targets.sort();
944 assert_eq!(targets, vec![b"AlphaError".to_vec(), b"BetaError".to_vec()]);
945 }
946
947 #[test]
948 fn extract_distinct_builtin_union_does_not_invent_receiver_type() {
949 let src = "def normalize(value: str | bytes):\n return value.strip()\n";
950 let g = extract_default("m.py", src, &make_anchor(), false);
951 let r = g
952 .refs()
953 .find(|r| {
954 r.kind == b"method_call"
955 && r.target.as_view().segments().last().unwrap().name == b"strip"
956 })
957 .expect("method_call strip");
958 assert_eq!(r.confidence, b"unresolved".to_vec());
959 }
960
961 #[test]
962 fn extract_workspace_union_parameter_emits_each_local_type_fact() {
963 let src = "class Alpha:\n pass\n\nclass Beta:\n pass\n\ndef render(value: Alpha | Beta):\n return value\n";
964 let g = extract_default("m.py", src, &make_anchor(), false);
965 let mut targets = g
966 .refs()
967 .filter(|reference| reference.kind == b"typed_as" && reference.alias == b"value")
968 .map(|reference| {
969 reference
970 .target
971 .as_view()
972 .segments()
973 .last()
974 .unwrap()
975 .name
976 .to_vec()
977 })
978 .collect::<Vec<_>>();
979 targets.sort();
980 assert_eq!(targets, vec![b"Alpha".to_vec(), b"Beta".to_vec()]);
981 }
982
983 #[test]
984 fn extract_optional_builtin_does_not_invent_receiver_type() {
985 let src = "from typing import Optional\n\ndef normalize(value: Optional[str]):\n return value.strip()\n";
986 let g = extract_default("m.py", src, &make_anchor(), false);
987 let r = g
988 .refs()
989 .find(|r| {
990 r.kind == b"method_call"
991 && r.target.as_view().segments().last().unwrap().name == b"strip"
992 })
993 .expect("method_call strip");
994 assert_eq!(r.confidence, b"unresolved".to_vec());
995 }
996
997 #[test]
998 fn extract_workspace_type_shadows_builtin_name() {
999 let src = "class str:\n def custom(self):\n pass\n\ndef use(value: str):\n value.custom()\n";
1000 let g = extract_default("m.py", src, &make_anchor(), false);
1001 let r = g
1002 .refs()
1003 .find(|r| {
1004 r.kind == b"method_call"
1005 && r.target.as_view().segments().last().unwrap().name == b"custom()"
1006 })
1007 .expect("method_call workspace str.custom");
1008 assert_eq!(r.confidence, b"resolved".to_vec());
1009 assert_eq!(r.target.parent().unwrap().parent(), Some(g.root().clone()));
1010 }
1011
1012 #[test]
1013 fn extract_literal_assignment_infers_builtin_receiver_type() {
1014 let src = "def normalize():\n value = ' user '\n return value.strip()\n";
1015 let g = extract_default("m.py", src, &make_anchor(), false);
1016 let r = g
1017 .refs()
1018 .find(|r| {
1019 r.kind == b"method_call"
1020 && r.target.as_view().segments().last().unwrap().name == b"strip"
1021 })
1022 .expect("method_call str.strip");
1023 assert_eq!(r.confidence, b"external".to_vec());
1024 assert_eq!(r.target.as_view().segments().nth(2).unwrap().name, b"str");
1025 }
1026
1027 #[test]
1028 fn extract_raw_bytes_literal_infers_bytes_receiver_type() {
1029 let src = "def decode():\n value = rb'user'\n return value.decode()\n";
1030 let g = extract_default("m.py", src, &make_anchor(), false);
1031 let r = g
1032 .refs()
1033 .find(|r| {
1034 r.kind == b"method_call"
1035 && r.target.as_view().segments().last().unwrap().name == b"decode"
1036 })
1037 .expect("method_call bytes.decode");
1038 assert_eq!(r.confidence, b"external".to_vec());
1039 assert_eq!(r.target.as_view().segments().nth(2).unwrap().name, b"bytes");
1040 }
1041
1042 #[test]
1043 fn extract_builtin_return_annotation_emits_external_returns_type() {
1044 let src = "def label() -> str:\n return 'user'\n";
1045 let g = extract_default("m.py", src, &make_anchor(), false);
1046 let r = g
1047 .refs()
1048 .find(|r| r.kind == b"returns_type")
1049 .expect("returns_type str");
1050 assert_eq!(r.confidence, b"external".to_vec());
1051 assert_eq!(r.target.as_view().segments().nth(2).unwrap().name, b"str");
1052 }
1053
1054 #[test]
1055 fn extract_async_return_type_requires_await_for_local_inference() {
1056 let src = "class User:\n def label(self) -> str:\n return 'user'\n\nasync def make() -> User:\n return User()\n\nasync def direct():\n value = make()\n return value.label()\n\nasync def awaited():\n value = await make()\n return value.label()\n";
1057 let g = extract_default("m.py", src, &make_anchor(), false);
1058 assert!(!g.refs().any(|r| r.kind == b"returns_type" && {
1059 g.defs().nth(r.source).is_some_and(|source| {
1060 source
1061 .moniker
1062 .as_view()
1063 .segments()
1064 .last()
1065 .is_some_and(|segment| segment.name == b"make()")
1066 })
1067 }));
1068 let method_calls = g
1069 .refs()
1070 .filter(|r| r.kind == b"method_call")
1071 .map(|r| {
1072 let source = g.defs().nth(r.source).unwrap();
1073 (
1074 source
1075 .moniker
1076 .as_view()
1077 .segments()
1078 .last()
1079 .unwrap()
1080 .name
1081 .to_vec(),
1082 r.confidence.to_vec(),
1083 )
1084 })
1085 .collect::<Vec<_>>();
1086 assert!(method_calls.contains(&(b"direct()".to_vec(), b"unresolved".to_vec())));
1087 assert!(method_calls.contains(&(b"awaited()".to_vec(), b"resolved".to_vec())));
1088 }
1089
1090 #[test]
1091 fn extract_async_method_does_not_emit_direct_return_type() {
1092 let src = "class User:\n pass\n\nclass Factory:\n async def make(self) -> User:\n return User()\n";
1093 let g = extract_default("m.py", src, &make_anchor(), false);
1094 assert!(!g.refs().any(|r| r.kind == b"returns_type"));
1095 }
1096
1097 #[test]
1098 fn extract_conflicting_factory_returns_do_not_type_local_assignment() {
1099 let src = "class First:\n def marker(self):\n pass\n\nclass Second:\n def marker(self):\n pass\n\ndef make() -> First:\n return First()\n\ndef make() -> Second:\n return Second()\n\ndef render():\n value = make()\n return value.marker()\n";
1100 let g = extract_default("m.py", src, &make_anchor(), false);
1101 let r = g
1102 .refs()
1103 .find(|r| {
1104 r.kind == b"method_call"
1105 && r.target.as_view().segments().last().unwrap().name == b"marker"
1106 })
1107 .expect("method_call marker");
1108 assert_eq!(r.confidence, b"unresolved".to_vec());
1109 assert!(!g.refs().any(|r| r.kind == b"returns_type"));
1110 }
1111
1112 #[test]
1113 fn extract_local_import_shadows_module_factory_return() {
1114 let src = "class Local:\n def marker(self):\n pass\n\ndef make() -> Local:\n return Local()\n\ndef render():\n from other import make\n value = make()\n return value.marker()\n";
1115 let g = extract_default("m.py", src, &make_anchor(), false);
1116 let r = g
1117 .refs()
1118 .find(|r| {
1119 r.kind == b"method_call"
1120 && r.target.as_view().segments().last().unwrap().name == b"marker"
1121 })
1122 .expect("method_call marker");
1123 assert_ne!(
1124 r.target
1125 .parent()
1126 .unwrap()
1127 .as_view()
1128 .segments()
1129 .last()
1130 .unwrap()
1131 .name,
1132 b"Local"
1133 );
1134 }
1135
1136 #[test]
1137 fn extract_aliased_optional_does_not_become_receiver_type() {
1138 let src = "from typing import Optional as Maybe\n\ndef normalize(value: Maybe[str]):\n return value.strip()\n";
1139 let g = extract_default("m.py", src, &make_anchor(), false);
1140 let r = g
1141 .refs()
1142 .find(|r| {
1143 r.kind == b"method_call"
1144 && r.target.as_view().segments().last().unwrap().name == b"strip"
1145 })
1146 .expect("method_call strip");
1147 assert_eq!(r.confidence, b"unresolved".to_vec());
1148 }
1149
1150 #[test]
1151 fn extract_typing_any_remains_dynamic() {
1152 let src = "from typing import Any\n\ndef make() -> Any:\n raise RuntimeError\n\ndef use(value: Any):\n return value.dynamic()\n";
1153 let g = extract_default("m.py", src, &make_anchor(), false);
1154 assert!(!g.refs().any(|r| r.kind == b"returns_type"));
1155 let r = g
1156 .refs()
1157 .find(|r| {
1158 r.kind == b"method_call"
1159 && r.target.as_view().segments().last().unwrap().name == b"dynamic"
1160 })
1161 .expect("method_call dynamic");
1162 assert_eq!(r.confidence, b"unresolved".to_vec());
1163 }
1164
1165 #[test]
1166 fn extract_conflicting_local_assignment_types_remain_ambiguous() {
1167 let src = "def mutate(flag):\n if flag:\n value = []\n else:\n value = {}\n value.append(1)\n";
1168 let g = extract_default("m.py", src, &make_anchor(), false);
1169 let r = g
1170 .refs()
1171 .find(|r| {
1172 r.kind == b"method_call"
1173 && r.target.as_view().segments().last().unwrap().name == b"append"
1174 })
1175 .expect("method_call append");
1176 assert_eq!(r.confidence, b"unresolved".to_vec());
1177 }
1178
1179 #[test]
1180 fn extract_unknown_reassignment_invalidates_known_local_type() {
1181 let src = "def mutate():\n value = []\n value = unknown()\n value.append(1)\n";
1182 let g = extract_default("m.py", src, &make_anchor(), false);
1183 let r = g
1184 .refs()
1185 .find(|r| {
1186 r.kind == b"method_call"
1187 && r.target.as_view().segments().last().unwrap().name == b"append"
1188 })
1189 .expect("method_call append");
1190 assert_eq!(r.confidence, b"unresolved".to_vec());
1191 }
1192
1193 #[test]
1194 fn extract_conflicting_instance_attribute_types_remain_ambiguous() {
1195 let src = "class Holder:\n def mutate(self, flag):\n if flag:\n self.value = []\n else:\n self.value = {}\n self.value.append(1)\n";
1196 let g = extract_default("m.py", src, &make_anchor(), false);
1197 let r = g
1198 .refs()
1199 .find(|r| {
1200 r.kind == b"method_call"
1201 && r.target.as_view().segments().last().unwrap().name == b"append"
1202 })
1203 .expect("method_call append");
1204 assert_eq!(r.confidence, b"unresolved".to_vec());
1205 }
1206
1207 #[test]
1208 fn extract_qualified_typing_any_does_not_bind_workspace_homonym() {
1209 let src = "import typing\n\nclass Any:\n def dynamic(self):\n pass\n\ndef use(value: typing.Any):\n value.dynamic()\n";
1210 let g = extract_default("m.py", src, &make_anchor(), false);
1211 let r = g
1212 .refs()
1213 .find(|r| {
1214 r.kind == b"method_call"
1215 && r.target.as_view().segments().last().unwrap().name == b"dynamic"
1216 })
1217 .expect("method_call dynamic");
1218 assert_eq!(r.confidence, b"unresolved".to_vec());
1219 }
1220
1221 #[test]
1222 fn extract_match_capture_shadows_module_factory() {
1223 let src = "class Local:\n def marker(self):\n pass\n\ndef make() -> Local:\n return Local()\n\ndef render(subject):\n value = make()\n match subject:\n case make:\n pass\n return value.marker()\n";
1224 let g = extract_default("m.py", src, &make_anchor(), false);
1225 let r = g
1226 .refs()
1227 .find(|r| {
1228 r.kind == b"method_call"
1229 && r.target.as_view().segments().last().unwrap().name == b"marker"
1230 })
1231 .expect("method_call marker");
1232 assert_eq!(r.confidence, b"unresolved".to_vec());
1233 }
1234
1235 #[test]
1236 fn extract_comprehension_walrus_shadows_module_factory() {
1237 let src = "class Local:\n def marker(self):\n pass\n\ndef make() -> Local:\n return Local()\n\ndef render(items):\n value = make()\n selected = [item for item in items if (make := item)]\n return value.marker()\n";
1238 let g = extract_default("m.py", src, &make_anchor(), false);
1239 let r = g
1240 .refs()
1241 .find(|r| {
1242 r.kind == b"method_call"
1243 && r.target.as_view().segments().last().unwrap().name == b"marker"
1244 })
1245 .expect("method_call marker");
1246 assert_eq!(r.confidence, b"unresolved".to_vec());
1247 }
1248
1249 #[test]
1250 fn extract_inner_param_shadows_outer_inferred_type() {
1251 let src =
1252 "def outer():\n value = []\n\n def inner(value):\n value.append(1)\n";
1253 let g = extract_default("m.py", src, &make_anchor(), false);
1254 let r = g
1255 .refs()
1256 .find(|r| {
1257 r.kind == b"method_call"
1258 && r.target.as_view().segments().last().unwrap().name == b"append"
1259 })
1260 .expect("method_call append");
1261 assert_eq!(r.confidence, b"unresolved".to_vec());
1262 }
1263
1264 #[test]
1265 fn extract_open_union_return_keeps_known_type_with_open_marker() {
1266 let src = "class User:\n pass\n\ndef maybe_make() -> User | None:\n return None\n";
1267 let g = extract_default("m.py", src, &make_anchor(), false);
1268 let fact = g
1269 .refs()
1270 .find(|reference| reference.kind == b"returns_type")
1271 .expect("known User return candidate");
1272 assert_eq!(
1273 fact.target.as_view().segments().last().unwrap().name,
1274 b"User"
1275 );
1276 assert_eq!(fact.receiver_hint, b"python_open_type_set".to_vec());
1277 }
1278
1279 #[test]
1280 fn extract_closed_union_return_emits_each_type_fact() {
1281 let src = "class Alpha:\n pass\n\nclass Beta:\n pass\n\ndef make() -> Alpha | Beta:\n return Alpha()\n";
1282 let g = extract_default("m.py", src, &make_anchor(), false);
1283 let mut targets = g
1284 .refs()
1285 .filter(|reference| reference.kind == b"returns_type")
1286 .map(|reference| {
1287 reference
1288 .target
1289 .as_view()
1290 .segments()
1291 .last()
1292 .unwrap()
1293 .name
1294 .to_vec()
1295 })
1296 .collect::<Vec<_>>();
1297 targets.sort();
1298 assert_eq!(targets, vec![b"Alpha".to_vec(), b"Beta".to_vec()]);
1299 }
1300
1301 #[test]
1302 fn extract_local_class_call_prefers_function_scoped_type() {
1303 let src = "class Local:\n pass\n\ndef make():\n class Local:\n pass\n return Local()\n";
1304 let g = extract_default("m.py", src, &make_anchor(), false);
1305 let r = g
1306 .refs()
1307 .find(|r| r.kind == b"instantiates")
1308 .expect("instantiates local Local");
1309 let target_last = r.target.as_view().segments().last().unwrap();
1310 assert_eq!(target_last.kind, b"class");
1311 assert_eq!(target_last.name, b"Local");
1312 let parent = r.target.parent().expect("class parent");
1313 let parent_last = parent.as_view().segments().last().unwrap();
1314 assert_eq!(parent_last.kind, b"function");
1315 assert_eq!(parent_last.name, b"make()");
1316 }
1317
1318 #[test]
1319 fn extract_keyword_argument_names_are_not_reads() {
1320 let src = "def save(value):\n return dict(id=value)\n";
1321 let g = extract_default("m.py", src, &make_anchor(), false);
1322 assert!(!g.refs().any(|r| {
1323 r.kind == b"reads" && r.target.as_view().segments().last().unwrap().name == b"id"
1324 }));
1325 }
1326
1327 #[test]
1328 fn extract_attribute_tail_is_not_a_bare_read() {
1329 let src = "def f(payment):\n return payment.id\n";
1330 let g = extract_default("m.py", src, &make_anchor(), false);
1331 assert!(!g.refs().any(|r| {
1332 r.kind == b"reads" && r.target.as_view().segments().last().unwrap().name == b"id"
1333 }));
1334 }
1335
1336 #[test]
1337 fn extract_function_docstring_emits_comment_def_parented_on_function() {
1338 let src = "def f():\n \"\"\"docstring\"\"\"\n return 0\n";
1339 let g = extract_default("m.py", src, &make_anchor(), false);
1340 let fn_moniker = MonikerBuilder::new()
1341 .project(b"app")
1342 .segment(b"lang", b"python")
1343 .segment(b"module", b"m")
1344 .segment(b"function", b"f()")
1345 .build();
1346 let docstring_count = g
1347 .defs()
1348 .filter(|d| d.kind == b"comment")
1349 .filter(|d| {
1350 d.parent
1351 .and_then(|i| g.defs().nth(i))
1352 .is_some_and(|p| p.moniker == fn_moniker)
1353 })
1354 .count();
1355 assert_eq!(
1356 docstring_count,
1357 1,
1358 "function docstring must emit one comment def parented on the function. defs: {:?}",
1359 g.def_monikers()
1360 );
1361 }
1362
1363 #[test]
1364 fn extract_class_and_module_docstring_each_emit_one_comment() {
1365 let src = "\"\"\"module doc\"\"\"\nclass A:\n \"\"\"class doc\"\"\"\n pass\n";
1366 let g = extract_default("m.py", src, &make_anchor(), false);
1367 assert_eq!(
1368 g.defs().filter(|d| d.kind == b"comment").count(),
1369 2,
1370 "module-level and class docstrings should each yield one comment def. defs: {:?}",
1371 g.def_monikers()
1372 );
1373 }
1374
1375 #[test]
1376 fn extract_non_docstring_string_at_start_is_not_a_comment() {
1377 let src = "x = \"hello\"\n";
1378 let g = extract_default("m.py", src, &make_anchor(), false);
1379 assert_eq!(
1380 g.defs().filter(|d| d.kind == b"comment").count(),
1381 0,
1382 "string literals that aren't bare expression-statement-strings must NOT be treated as docstrings"
1383 );
1384 }
1385}