use tree_sitter::{Language, Parser, Tree};
use crate::core::code_graph::CodeGraph;
use crate::core::moniker::Moniker;
use crate::core::shape::Shape;
use crate::lang::KindSpec;
pub mod build;
mod canonicalize;
mod kinds;
mod sdk_pipeline;
#[derive(Clone, Debug, Default)]
pub struct Presets {}
pub fn parse(source: &str) -> Tree {
let mut parser = Parser::new();
let language: Language = tree_sitter_python::LANGUAGE.into();
parser.set_language(&language).unwrap_or_else(|err| {
panic!("failed to load tree-sitter Python grammar: {err}");
});
parser.parse(source, None).unwrap_or_else(|| {
panic!("tree-sitter parse returned None on a non-cancelled call");
})
}
pub fn extract(
uri: &str,
source: &str,
anchor: &Moniker,
deep: bool,
presets: &Presets,
) -> CodeGraph {
sdk_pipeline::extract(uri, source, anchor, deep, presets)
}
pub struct Lang;
const DEF_KINDS: &[&str] = &[
"class",
"type",
"function",
"method",
"async_function",
"path",
];
const DEF_KIND_SPECS: &[KindSpec] = &[
KindSpec::new("class", Shape::Type, 20, "class"),
KindSpec::new("type", Shape::Type, 21, "type"),
KindSpec::new("function", Shape::Callable, 40, "function"),
KindSpec::new("async_function", Shape::Callable, 41, "async_function"),
KindSpec::new("method", Shape::Callable, 42, "method"),
KindSpec::new("path", Shape::Value, 63, "path"),
];
impl crate::lang::LangExtractor for Lang {
type Presets = Presets;
const LANG_TAG: &'static str = "python";
const ALLOWED_KINDS: &'static [&'static str] = DEF_KINDS;
const KIND_SPECS: &'static [KindSpec] = DEF_KIND_SPECS;
const ALLOWED_VISIBILITIES: &'static [&'static str] = &["public", "private", "module"];
fn extract(
uri: &str,
source: &str,
anchor: &Moniker,
deep: bool,
presets: &Self::Presets,
) -> CodeGraph {
extract(uri, source, anchor, deep, presets)
}
}
pub fn builtin_external_root(root: &str) -> bool {
root == "builtins" || sdk_pipeline::STDLIB_PACKAGES.binary_search(&root).is_ok()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::moniker::MonikerBuilder;
use crate::lang::assert_conformance;
fn make_anchor() -> Moniker {
MonikerBuilder::new().project(b"app").build()
}
fn extract_default(uri: &str, source: &str, anchor: &Moniker, deep: bool) -> CodeGraph {
let g = extract(uri, source, anchor, deep, &Presets::default());
assert_conformance::<super::Lang>(&g, anchor);
g
}
#[test]
fn parse_empty_returns_module() {
let tree = parse("");
assert_eq!(tree.root_node().kind(), "module");
}
#[test]
fn extract_module_uses_path_segments() {
let g = extract_default("acme/util/text.py", "", &make_anchor(), false);
let expected = MonikerBuilder::new()
.project(b"app")
.segment(b"lang", b"python")
.segment(b"package", b"acme")
.segment(b"package", b"util")
.segment(b"module", b"text")
.build();
assert_eq!(g.root(), &expected);
}
#[test]
fn extract_module_root_is_filename_only() {
let g = extract_default("foo.py", "", &make_anchor(), false);
let expected = MonikerBuilder::new()
.project(b"app")
.segment(b"lang", b"python")
.segment(b"module", b"foo")
.build();
assert_eq!(g.root(), &expected);
}
#[test]
fn extract_function_with_typed_params_emits_full_signature() {
let src = "def make(x: int, y: str) -> int:\n return x\n";
let g = extract_default("m.py", src, &make_anchor(), false);
let f = g
.defs()
.find(|d| d.kind == b"function")
.expect("function def");
let last = f.moniker.as_view().segments().last().unwrap();
assert_eq!(last.kind, b"function");
assert_eq!(last.name, b"make(x:int,y:str)");
assert_eq!(f.signature, b"x:int,y:str".to_vec());
}
#[test]
fn extract_function_with_untyped_params_uses_name_only_slots() {
let src = "def f(a, b=1):\n return a\n";
let g = extract_default("m.py", src, &make_anchor(), false);
let f = g
.defs()
.find(|d| d.kind == b"function")
.expect("function def");
let last = f.moniker.as_view().segments().last().unwrap();
assert_eq!(last.name, b"f(a,b)");
assert_eq!(f.signature, b"a,b".to_vec());
}
#[test]
fn extract_classmethod_excludes_cls_from_signature() {
let src = "class Foo:\n @classmethod\n def make(cls, x: int) -> 'Foo':\n return cls()\n";
let g = extract_default("foo.py", src, &make_anchor(), false);
let m = g.defs().find(|d| d.kind == b"method").expect("method def");
assert_eq!(
m.moniker.as_view().segments().last().unwrap().name,
b"make(x:int)"
);
}
#[test]
fn extract_double_underscore_visibility_is_private() {
let src = "class Foo:\n def __secret(self):\n pass\n";
let g = extract_default("foo.py", src, &make_anchor(), false);
let m = g.defs().find(|d| d.kind == b"method").expect("method def");
assert_eq!(m.visibility, b"private".to_vec());
}
#[test]
fn extract_single_underscore_visibility_is_module() {
let src = "def _internal():\n pass\n";
let g = extract_default("m.py", src, &make_anchor(), false);
let f = g
.defs()
.find(|d| d.kind == b"function")
.expect("function def");
assert_eq!(f.visibility, b"module".to_vec());
}
#[test]
fn extract_import_module_emits_imports_module() {
let src = "import os\nimport acme.util as u\n";
let g = extract_default("m.py", src, &make_anchor(), false);
let kinds: Vec<&[u8]> = g.refs().map(|r| r.kind.as_ref()).collect();
assert_eq!(kinds.iter().filter(|k| **k == b"imports_module").count(), 2);
}
#[test]
fn extract_stdlib_import_marks_external() {
let g = extract_default("m.py", "import json\n", &make_anchor(), false);
let r = g
.refs()
.find(|r| r.kind == b"imports_module")
.expect("imports_module");
assert_eq!(r.confidence, b"external".to_vec());
}
#[test]
fn extract_project_import_marks_imported() {
let g = extract_default("m.py", "import acme.util\n", &make_anchor(), false);
let r = g
.refs()
.find(|r| r.kind == b"imports_module")
.expect("imports_module");
assert_eq!(r.confidence, b"imported".to_vec());
}
#[test]
fn extract_from_import_emits_one_imports_symbol_per_name() {
let src = "from acme.util import a, b as c\n";
let g = extract_default("m.py", src, &make_anchor(), false);
let names: Vec<&[u8]> = g
.refs()
.filter(|r| r.kind == b"imports_symbol")
.map(|r| r.target.as_view().segments().last().unwrap().name)
.collect();
assert_eq!(names, vec![&b"a"[..], &b"b"[..]]);
let segs: Vec<_> = g
.refs()
.find(|r| r.kind == b"imports_symbol")
.unwrap()
.target
.as_view()
.segments()
.collect();
let kinds: Vec<&[u8]> = segs.iter().map(|s| s.kind).collect();
assert_eq!(
kinds,
vec![&b"lang"[..], &b"package"[..], &b"module"[..], &b"path"[..]]
);
let aliased = g
.refs()
.find(|r| r.kind == b"imports_symbol" && r.alias == b"c")
.expect("aliased import");
assert_eq!(aliased.alias, b"c".to_vec());
}
#[test]
fn extract_relative_import_resolves_against_importer() {
let src = "from .util import helper\n";
let g = extract_default("acme/m.py", src, &make_anchor(), false);
let r = g
.refs()
.find(|r| r.kind == b"imports_symbol")
.expect("imports_symbol");
let segs: Vec<_> = r.target.as_view().segments().collect();
let kinds: Vec<&[u8]> = segs.iter().map(|s| s.kind).collect();
let names: Vec<&[u8]> = segs.iter().map(|s| s.name).collect();
assert_eq!(
kinds,
vec![&b"lang"[..], &b"package"[..], &b"module"[..], &b"path"[..]]
);
assert_eq!(
names,
vec![&b"python"[..], &b"acme"[..], &b"util"[..], &b"helper"[..]]
);
}
#[test]
fn extract_relative_import_underflow_falls_back_to_external_pkg() {
let src = "from ...foo import bar\n";
let g = extract_default("m.py", src, &make_anchor(), false);
let r = g
.refs()
.find(|r| r.kind == b"imports_symbol")
.expect("imports_symbol");
let segs: Vec<_> = r.target.as_view().segments().collect();
assert_eq!(segs[0].kind, b"external_pkg");
assert_eq!(segs[0].name, b"...");
}
#[test]
fn extract_decorator_emits_annotates() {
let src = "import functools\n@functools.wraps(fn)\ndef g():\n pass\n";
let g = extract_default("m.py", src, &make_anchor(), false);
let ann = g
.refs()
.find(|r| r.kind == b"annotates")
.expect("annotates ref");
assert_eq!(
ann.target.as_view().segments().last().unwrap().name,
b"wraps"
);
}
#[test]
fn extract_param_read_marks_confidence_local() {
let src = "def f(x):\n return x\n";
let g = extract_default("m.py", src, &make_anchor(), true);
let r = g
.refs()
.find(|r| {
r.kind == b"reads" && r.target.as_view().segments().last().unwrap().name == b"x"
})
.expect("reads x");
assert_eq!(r.confidence, b"local".to_vec());
}
#[test]
fn extract_deep_emits_param_def() {
let src = "def f(x, y: int):\n return x\n";
let g = extract_default("m.py", src, &make_anchor(), true);
let params: Vec<&[u8]> = g
.defs()
.filter(|d| d.kind == b"param")
.map(|d| d.moniker.as_view().segments().last().unwrap().name)
.collect();
assert!(params.contains(&&b"x"[..]));
assert!(params.contains(&&b"y"[..]));
}
#[test]
fn extract_async_function_tracks_local_assignments() {
let src = "async def f():\n x = 1\n return x\n";
let g = extract_default("m.py", src, &make_anchor(), true);
let r = g
.refs()
.find(|r| {
r.kind == b"reads" && r.target.as_view().segments().last().unwrap().name == b"x"
})
.expect("reads x");
assert_eq!(r.confidence, b"local".to_vec());
}
#[test]
fn extract_imported_identifier_read_targets_import_moniker() {
let src = "import asyncio\nasync def f():\n return asyncio\n";
let g = extract_default("m.py", src, &make_anchor(), false);
let r = g
.refs()
.find(|r| r.kind == b"reads" && r.confidence == b"external")
.expect("external read");
let segs: Vec<_> = r.target.as_view().segments().collect();
assert_eq!(segs[0].kind, b"external_pkg");
assert_eq!(segs[0].name, b"asyncio");
}
#[test]
fn extract_self_member_call_resolves_via_typed_constructor_param() {
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";
let g = extract_default("m.py", src, &make_anchor(), false);
let r = g
.refs()
.find(|r| {
r.kind == b"method_call"
&& r.target.as_view().segments().last().unwrap().name == b"reserve()"
})
.expect("method_call Store.reserve");
assert_eq!(r.confidence, b"resolved".to_vec());
let parent = r.target.parent().expect("method parent");
assert_eq!(parent.as_view().segments().last().unwrap().name, b"Store");
}
#[test]
fn extract_self_member_call_uses_constructor_type_independent_of_method_order() {
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";
let g = extract_default("m.py", src, &make_anchor(), false);
let targets: Vec<Moniker> = g
.refs()
.filter(|r| {
r.kind == b"method_call"
&& r.target.as_view().segments().last().unwrap().name == b"reserve()"
})
.map(|r| r.target.clone())
.collect();
assert_eq!(targets.len(), 2);
for target in targets {
let parent = target.parent().expect("method parent");
assert_eq!(parent.as_view().segments().last().unwrap().name, b"Store");
}
}
#[test]
fn extract_self_callable_attr_call_targets_attribute_type_alias() {
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";
let g = extract_default("m.py", src, &make_anchor(), false);
let r = g
.refs()
.find(|r| {
r.kind == b"calls"
&& r.target.as_view().segments().last().unwrap().name == b"Callback"
})
.expect("calls Callback type alias");
assert_eq!(r.confidence, b"resolved".to_vec());
assert!(!g.refs().any(|r| {
r.kind == b"method_call" && r.target.as_view().segments().last().unwrap().name == b"_cb"
}));
}
#[test]
fn extract_type_alias_emits_type_def_and_rhs_uses_type() {
let src = "from typing import TypeAlias\n\nclass User:\n pass\n\nUserMap: TypeAlias = dict[str, User]\ntype UserResult = User | None\n";
let g = extract_default("m.py", src, &make_anchor(), false);
let type_names: Vec<&[u8]> = g
.defs()
.filter(|d| d.kind == b"type")
.map(|d| d.moniker.as_view().segments().last().unwrap().name)
.collect();
assert!(type_names.contains(&&b"UserMap"[..]));
assert!(type_names.contains(&&b"UserResult"[..]));
assert!(g.refs().any(|r| {
r.kind == b"uses_type"
&& r.target.as_view().segments().last().unwrap().name == b"User"
&& r.confidence == b"resolved"
}));
assert!(!g.refs().any(|r| {
r.kind == b"uses_type"
&& matches!(
r.target.as_view().segments().last().unwrap().name,
b"dict" | b"str" | b"None"
)
}));
}
#[test]
fn extract_local_class_call_emits_instantiates() {
let src = "class User:\n pass\n\ndef make() -> User:\n return User()\n";
let g = extract_default("m.py", src, &make_anchor(), false);
let r = g
.refs()
.find(|r| r.kind == b"instantiates")
.expect("instantiates User");
assert_eq!(r.confidence, b"resolved".to_vec());
assert_eq!(r.target.as_view().segments().last().unwrap().name, b"User");
}
#[test]
fn extract_local_class_call_prefers_function_scoped_type() {
let src = "class Local:\n pass\n\ndef make():\n class Local:\n pass\n return Local()\n";
let g = extract_default("m.py", src, &make_anchor(), false);
let r = g
.refs()
.find(|r| r.kind == b"instantiates")
.expect("instantiates local Local");
let target_last = r.target.as_view().segments().last().unwrap();
assert_eq!(target_last.kind, b"class");
assert_eq!(target_last.name, b"Local");
let parent = r.target.parent().expect("class parent");
let parent_last = parent.as_view().segments().last().unwrap();
assert_eq!(parent_last.kind, b"function");
assert_eq!(parent_last.name, b"make()");
}
#[test]
fn extract_keyword_argument_names_are_not_reads() {
let src = "def save(value):\n return dict(id=value)\n";
let g = extract_default("m.py", src, &make_anchor(), false);
assert!(!g.refs().any(|r| {
r.kind == b"reads" && r.target.as_view().segments().last().unwrap().name == b"id"
}));
}
#[test]
fn extract_attribute_tail_is_not_a_bare_read() {
let src = "def f(payment):\n return payment.id\n";
let g = extract_default("m.py", src, &make_anchor(), false);
assert!(!g.refs().any(|r| {
r.kind == b"reads" && r.target.as_view().segments().last().unwrap().name == b"id"
}));
}
#[test]
fn extract_function_docstring_emits_comment_def_parented_on_function() {
let src = "def f():\n \"\"\"docstring\"\"\"\n return 0\n";
let g = extract_default("m.py", src, &make_anchor(), false);
let fn_moniker = MonikerBuilder::new()
.project(b"app")
.segment(b"lang", b"python")
.segment(b"module", b"m")
.segment(b"function", b"f()")
.build();
let docstring_count = g
.defs()
.filter(|d| d.kind == b"comment")
.filter(|d| {
d.parent
.and_then(|i| g.defs().nth(i))
.is_some_and(|p| p.moniker == fn_moniker)
})
.count();
assert_eq!(
docstring_count,
1,
"function docstring must emit one comment def parented on the function. defs: {:?}",
g.def_monikers()
);
}
#[test]
fn extract_class_and_module_docstring_each_emit_one_comment() {
let src = "\"\"\"module doc\"\"\"\nclass A:\n \"\"\"class doc\"\"\"\n pass\n";
let g = extract_default("m.py", src, &make_anchor(), false);
assert_eq!(
g.defs().filter(|d| d.kind == b"comment").count(),
2,
"module-level and class docstrings should each yield one comment def. defs: {:?}",
g.def_monikers()
);
}
#[test]
fn extract_non_docstring_string_at_start_is_not_a_comment() {
let src = "x = \"hello\"\n";
let g = extract_default("m.py", src, &make_anchor(), false);
assert_eq!(
g.defs().filter(|d| d.kind == b"comment").count(),
0,
"string literals that aren't bare expression-statement-strings must NOT be treated as docstrings"
);
}
}