use std::path::{Path, PathBuf};
use rayon::prelude::*;
use crate::error::{Diagnostic, Severity};
use crate::language::LangId;
use crate::model::{IdGenerator, Symbol, SymbolId};
use crate::parser;
pub use crate::model::FileExtraction;
#[derive(Debug, Clone, Copy, Default)]
pub struct ExtractOptions {
pub skip_imports_and_refs: bool,
}
pub struct ExtractionResult {
pub files: Vec<FileExtraction>,
}
#[derive(Debug, Default)]
pub struct ExtractionIdGenerators {
symbols: IdGenerator<SymbolId>,
#[cfg(feature = "dataflow")]
data_nodes: IdGenerator<crate::model::DataNodeId>,
}
impl ExtractionIdGenerators {
pub fn new() -> Self {
Self::default()
}
pub fn with_symbol_start(symbol_start: u32) -> Self {
Self {
symbols: IdGenerator::with_start(symbol_start),
#[cfg(feature = "dataflow")]
data_nodes: IdGenerator::new(),
}
}
#[cfg(feature = "dataflow")]
pub fn with_starts(symbol_start: u32, data_node_start: u32) -> Self {
Self {
symbols: IdGenerator::with_start(symbol_start),
data_nodes: IdGenerator::with_start(data_node_start),
}
}
pub fn symbols(&self) -> &IdGenerator<SymbolId> {
&self.symbols
}
#[cfg(feature = "dataflow")]
pub fn data_nodes(&self) -> &IdGenerator<crate::model::DataNodeId> {
&self.data_nodes
}
}
#[derive(Debug, Clone, Copy)]
pub struct InMemorySource<'a> {
pub uri: &'a str,
pub text: &'a str,
pub version: i32,
pub language: LangId,
}
#[derive(Debug, Clone)]
pub struct VersionedExtraction {
pub uri: String,
pub version: i32,
pub file: FileExtraction,
}
pub fn extract(files: &[(std::path::PathBuf, LangId)]) -> ExtractionResult {
extract_with_options(files, &ExtractOptions::default())
}
pub fn extract_with_options(
files: &[(std::path::PathBuf, LangId)],
opts: &ExtractOptions,
) -> ExtractionResult {
let id_generators = ExtractionIdGenerators::new();
extract_with_id_gen(files, opts, &id_generators)
}
pub fn extract_with_id_gen(
files: &[(PathBuf, LangId)],
opts: &ExtractOptions,
id_generators: &ExtractionIdGenerators,
) -> ExtractionResult {
let mut file_extractions: Vec<_> = files
.par_iter()
.map(|(path, lang)| extract_single_file(path, lang, id_generators, opts))
.collect();
file_extractions.sort_by(|a, b| a.path.cmp(&b.path));
ExtractionResult {
files: file_extractions,
}
}
fn extract_single_file(
path: &Path,
lang: &LangId,
id_generators: &ExtractionIdGenerators,
opts: &ExtractOptions,
) -> FileExtraction {
let source = match std::fs::read(path) {
Ok(source) => source,
Err(error) => {
return failed_extraction(path, *lang, format!("failed to read file: {error}"));
}
};
extract_source(path, *lang, &source, id_generators, opts)
}
pub fn extract_text_with_id_gen(
source: InMemorySource<'_>,
opts: &ExtractOptions,
id_generators: &ExtractionIdGenerators,
) -> Result<VersionedExtraction, crate::Error> {
let parsed_uri =
url::Url::parse(source.uri).map_err(|error| crate::Error::InvalidSourceUri {
uri: source.uri.to_string(),
message: error.to_string(),
})?;
let path = parsed_uri
.to_file_path()
.map_err(|()| crate::Error::InvalidSourceUri {
uri: source.uri.to_string(),
message: "URI must use the file scheme and contain an absolute path".to_string(),
})?;
let path = crate::input::simplified_path(&path).to_path_buf();
let file = extract_source(
&path,
source.language,
source.text.as_bytes(),
id_generators,
opts,
);
Ok(VersionedExtraction {
uri: source.uri.to_string(),
version: source.version,
file,
})
}
fn extract_source(
path: &Path,
lang: LangId,
source: &[u8],
id_generators: &ExtractionIdGenerators,
opts: &ExtractOptions,
) -> FileExtraction {
let tree = match crate::parser::parse_tree(lang, source) {
Ok(t) => t,
Err(e) => {
return failed_extraction(path, lang, e.to_string());
}
};
let metrics = parser::tree_metrics(&tree, source);
let mut diags = Vec::new();
if metrics.error_ratio > 0.5 {
diags.push(Diagnostic {
path: path.to_path_buf(),
severity: Severity::Warning,
message: format!(
"file has {:.0}% parse errors, results may be incomplete",
metrics.error_ratio * 100.0
),
source_range: None,
});
}
let raw_symbols = crate::language::extract_symbols_for(lang, &tree, source);
let symbols = raw_symbols
.into_iter()
.map(|raw| Symbol {
id: id_generators.symbols.next(),
name: raw.name.into_owned(),
kind: raw.kind,
language: lang,
file_path: path.to_path_buf(),
source_range: raw.source_range,
visibility: raw.visibility,
signature: raw.signature.map(|s| s.into_owned()),
docstring: raw.docstring.map(|s| s.into_owned()),
is_async: raw.is_async,
})
.collect();
let (imports, references) = if opts.skip_imports_and_refs {
(Vec::new(), Vec::new())
} else {
crate::language::extract_imports_and_references_for(lang, &tree, source, path)
};
#[cfg(feature = "metacall-deploy")]
let call_sites = crate::deploy::scanner::scan_file(lang, &tree, source, path);
#[cfg(feature = "dataflow")]
let (data_nodes, flow_edges) =
crate::language::dataflow::extract_dataflow(lang, &tree, source, &id_generators.data_nodes);
let mut out = FileExtraction::empty(path.to_path_buf(), lang);
out.symbols = symbols;
out.imports = imports;
out.references = references;
out.diagnostics = diags;
out.ast_node_count = metrics.node_count;
#[cfg(feature = "metacall-deploy")]
{
out.call_sites = call_sites;
}
#[cfg(feature = "dataflow")]
{
out.data_nodes = data_nodes;
out.flow_edges = flow_edges;
}
out
}
fn failed_extraction(path: &Path, lang: LangId, message: String) -> FileExtraction {
FileExtraction::failed(path.to_path_buf(), lang, message)
}
#[cfg(test)]
mod tests {
use super::*;
use std::path::PathBuf;
fn test_dir() -> PathBuf {
let dir = std::env::temp_dir().join("meta_ast_test_extractor");
let _ = std::fs::create_dir_all(&dir);
dir
}
fn write_temp(name: &str, content: &[u8]) -> PathBuf {
let path = test_dir().join(name);
std::fs::write(&path, content).unwrap();
path
}
#[test]
fn extract_single_python_file() {
let path = write_temp("single.py", b"def hello(): pass\n");
let result = extract(&[(path.clone(), LangId::Python)]);
assert_eq!(result.files.len(), 1);
assert!(!result.files[0].symbols.is_empty());
assert!(result.files[0].diagnostics.is_empty());
let names: Vec<&str> = result.files[0]
.symbols
.iter()
.map(|s| s.name.as_str())
.collect();
assert!(names.contains(&"hello"));
}
#[test]
fn extract_multiple_files_parallel() {
let p1 = write_temp("file_a.py", b"def alpha(): pass\n");
let p2 = write_temp("file_b.py", b"def beta(): pass\ndef gamma(): pass\n");
let p3 = write_temp("file_c.py", b"class Delta: pass\n");
let files = vec![
(p1.clone(), LangId::Python),
(p2.clone(), LangId::Python),
(p3.clone(), LangId::Python),
];
let result = extract(&files);
let all_names: Vec<&str> = result
.files
.iter()
.flat_map(|f| f.symbols.iter().map(|s| s.name.as_str()))
.collect();
assert!(all_names.contains(&"alpha"), "missing alpha: {all_names:?}");
assert!(all_names.contains(&"beta"), "missing beta: {all_names:?}");
assert!(all_names.contains(&"gamma"), "missing gamma: {all_names:?}");
assert!(all_names.contains(&"Delta"), "missing Delta: {all_names:?}");
}
#[test]
fn accumulate_diagnostics_on_malformed() {
let path = test_dir().join("nonexistent_broken.py");
let _ = std::fs::remove_file(&path);
let result = extract(&[(path, LangId::Python)]);
assert!(!result.files[0].diagnostics.is_empty());
}
#[test]
fn partial_extraction_on_errors() {
let valid = write_temp("valid_partial.py", b"def works(): pass\n");
let broken = write_temp(
"broken_partial.py",
b"def broken(\n # missing close paren and colon\n",
);
let result = extract(&[(valid.clone(), LangId::Python), (broken, LangId::Python)]);
let names: Vec<&str> = result
.files
.iter()
.flat_map(|f| f.symbols.iter().map(|s| s.name.as_str()))
.collect();
assert!(
names.contains(&"works"),
"valid file symbols should be present: {names:?}"
);
}
#[test]
fn output_deterministic() {
let path = write_temp("deterministic.py", b"def foo(): pass\ndef bar(): pass\n");
let files = vec![(path.clone(), LangId::Python)];
let r1 = extract(&files);
let r2 = extract(&files);
let names1: Vec<String> = r1
.files
.iter()
.flat_map(|f| f.symbols.iter().map(|s| s.name.clone()))
.collect();
let names2: Vec<String> = r2
.files
.iter()
.flat_map(|f| f.symbols.iter().map(|s| s.name.clone()))
.collect();
assert_eq!(names1, names2);
}
#[test]
fn in_memory_extraction_uses_unsaved_text_and_preserves_version() {
let path = test_dir().join("buffer.py");
let _ = std::fs::remove_file(&path);
let uri = url::Url::from_file_path(&path).unwrap().to_string();
let id_generators = ExtractionIdGenerators::with_symbol_start(40);
let result = extract_text_with_id_gen(
InMemorySource {
uri: &uri,
text: "def unsaved(): pass\n",
version: 7,
language: LangId::Python,
},
&ExtractOptions::default(),
&id_generators,
)
.unwrap();
assert_eq!(result.version, 7);
assert_eq!(result.uri, uri);
assert_eq!(result.file.path, path);
assert_eq!(result.file.symbols[0].name, "unsaved");
assert_eq!(result.file.symbols[0].id, SymbolId::new(40).unwrap());
}
#[test]
fn in_memory_extraction_rejects_non_file_uri() {
let id_generators = ExtractionIdGenerators::new();
let error = extract_text_with_id_gen(
InMemorySource {
uri: "untitled:buffer.py",
text: "def value(): pass\n",
version: 1,
language: LangId::Python,
},
&ExtractOptions::default(),
&id_generators,
)
.unwrap_err();
assert!(matches!(error, crate::Error::InvalidSourceUri { .. }));
}
#[test]
fn extraction_id_generators_accessors() {
let id_generators = ExtractionIdGenerators::with_symbol_start(100);
assert_eq!(id_generators.symbols().next(), SymbolId::new(100).unwrap());
#[cfg(feature = "dataflow")]
{
let dual = ExtractionIdGenerators::with_starts(200, 300);
assert_eq!(dual.symbols().next(), SymbolId::new(200).unwrap());
assert_eq!(
dual.data_nodes().next(),
crate::model::DataNodeId::new(300).unwrap()
);
}
}
#[test]
fn symbols_assigned_ids() {
let path = write_temp("ids.py", b"def a(): pass\ndef b(): pass\ndef c(): pass\n");
let result = extract(&[(path, LangId::Python)]);
let ids: Vec<u32> = result.files[0]
.symbols
.iter()
.map(|s| s.id.to_raw())
.collect();
let mut sorted_ids = ids.clone();
sorted_ids.sort();
assert_eq!(ids, sorted_ids, "IDs should be sequential");
for window in sorted_ids.windows(2) {
assert_eq!(window[1] - window[0], 1, "IDs should be consecutive");
}
let unique: std::collections::HashSet<u32> = ids.iter().copied().collect();
assert_eq!(
unique.len(),
result.files[0].symbols.len(),
"all IDs must be unique"
);
}
}