use std::path::{Path, PathBuf};
use crate::incremental::types::{DependencyEdge, DependencyType};
#[derive(Debug, thiserror::Error)]
pub enum ExtractionError {
#[error("parse error: failed to parse Go source")]
ParseError,
#[error("query error: {0}")]
QueryError(String),
#[error("unresolved import: {path}")]
UnresolvedImport {
path: String,
},
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ImportInfo {
pub import_path: String,
pub alias: Option<String>,
pub is_dot_import: bool,
pub is_blank_import: bool,
}
#[derive(Debug, Clone)]
pub struct GoDependencyExtractor {
module_path: Option<String>,
vendor_mode: bool,
}
impl GoDependencyExtractor {
pub fn new(module_path: Option<String>) -> Self {
Self {
module_path,
vendor_mode: false,
}
}
pub fn with_vendor(module_path: Option<String>, vendor_mode: bool) -> Self {
Self {
module_path,
vendor_mode,
}
}
pub fn extract_imports(
&self,
source: &str,
_file_path: &Path,
) -> Result<Vec<ImportInfo>, ExtractionError> {
if source.is_empty() {
return Ok(Vec::new());
}
let language = thread_language::parsers::language_go();
let mut parser = tree_sitter::Parser::new();
parser
.set_language(&language)
.map_err(|_| ExtractionError::ParseError)?;
let tree = parser
.parse(source, None)
.ok_or(ExtractionError::ParseError)?;
let root_node = tree.root_node();
let mut imports = Vec::new();
self.walk_imports(root_node, source.as_bytes(), &mut imports);
Ok(imports)
}
fn walk_imports(
&self,
node: tree_sitter::Node<'_>,
source: &[u8],
imports: &mut Vec<ImportInfo>,
) {
if node.kind() == "import_declaration" {
self.extract_from_import_declaration(node, source, imports);
return;
}
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
self.walk_imports(child, source, imports);
}
}
fn extract_from_import_declaration(
&self,
node: tree_sitter::Node<'_>,
source: &[u8],
imports: &mut Vec<ImportInfo>,
) {
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
match child.kind() {
"import_spec" => {
if let Some(info) = self.parse_import_spec(child, source) {
imports.push(info);
}
}
"import_spec_list" => {
let mut list_cursor = child.walk();
for spec in child.children(&mut list_cursor) {
if spec.kind() == "import_spec"
&& let Some(info) = self.parse_import_spec(spec, source)
{
imports.push(info);
}
}
}
_ => {}
}
}
}
fn parse_import_spec(&self, node: tree_sitter::Node<'_>, source: &[u8]) -> Option<ImportInfo> {
let mut alias: Option<String> = None;
let mut is_dot_import = false;
let mut is_blank_import = false;
let mut import_path: Option<String> = None;
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
match child.kind() {
"dot" => {
is_dot_import = true;
}
"blank_identifier" => {
is_blank_import = true;
}
"package_identifier" => {
let name = child.utf8_text(source).ok()?.to_string();
alias = Some(name);
}
"interpreted_string_literal" => {
let raw = child.utf8_text(source).ok()?;
let path = raw.trim_matches('"').to_string();
import_path = Some(path);
}
_ => {}
}
}
import_path.map(|path| ImportInfo {
import_path: path,
alias,
is_dot_import,
is_blank_import,
})
}
pub fn resolve_import_path(
&self,
_source_file: &Path,
import_path: &str,
) -> Result<PathBuf, ExtractionError> {
if let Some(ref module) = self.module_path
&& let Some(relative) = import_path.strip_prefix(module)
{
let relative = relative.strip_prefix('/').unwrap_or(relative);
return Ok(PathBuf::from(relative));
}
if self.vendor_mode {
return Ok(PathBuf::from(format!("vendor/{import_path}")));
}
Err(ExtractionError::UnresolvedImport {
path: import_path.to_string(),
})
}
pub fn extract_dependency_edges(
&self,
source: &str,
file_path: &Path,
) -> Result<Vec<DependencyEdge>, ExtractionError> {
let imports = self.extract_imports(source, file_path)?;
let mut edges = Vec::new();
for import in &imports {
if let Ok(resolved) = self.resolve_import_path(file_path, &import.import_path) {
edges.push(DependencyEdge::new(
file_path.to_path_buf(),
resolved,
DependencyType::Import,
));
}
}
Ok(edges)
}
}