use alloc::vec::Vec;
use crate::query::{adj_ast, match_ast_d, match_ast_group_d};
use crate::syntax::builders::file::build_file_node;
use crate::syntax::fields::python::import_from_statement;
use crate::syntax::{SyntaxGraphArgs, SyntaxGraphError};
use crate::{CodeGraph, NodeId};
fn imported_names(args: &SyntaxGraphArgs<'_>, import_id: NodeId) -> Vec<NodeId> {
let graph = args.ast_graph;
let mut names = match_ast_group_d(graph, import_id, "dotted_name", Some(-1));
if match_ast_d(graph, import_id, "wildcard_import", Some(1)).is_some() {
return names;
}
let module_n_id = match graph.label_type(import_id) {
Some("import_from_statement") => args
.required_field_alt(import_id, import_from_statement::MODULE_NAME)
.ok(),
_ => None,
};
let Some(module_n_id) = module_n_id else {
return names;
};
let drop_id = if names.contains(&module_n_id) {
Some(module_n_id)
} else {
match_ast_d(graph, module_n_id, "dotted_name", Some(1))
};
if let Some(pos) = drop_id.and_then(|id| names.iter().position(|name| *name == id)) {
names.remove(pos);
}
names
}
pub fn reader(
args: &mut SyntaxGraphArgs<'_>,
n_id: NodeId,
) -> Result<Option<NodeId>, SyntaxGraphError> {
let graph = args.ast_graph;
let mut filtered_ids: Vec<NodeId> = Vec::new();
for c_id in adj_ast(graph, n_id, Some(1), &[]) {
if matches!(
graph.label_type(c_id),
Some("import_statement" | "import_from_statement")
) {
filtered_ids.extend(imported_names(args, c_id));
} else {
filtered_ids.push(c_id);
}
}
build_file_node(args, n_id, &filtered_ids).map(Some)
}