Skip to main content

code_moniker_core/lang/python/
mod.rs

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}