archlens/crates/adapters/python-project/src/python_project_analyzer.rs

use std::collections::{HashMap, HashSet};
use std::path::Path;

use serde::Deserialize;

use archlens_domain::{
    CodeElement, CodeElementKind, CodeGraph, DomainError, FilePath, Relationship, RelationshipKind,
    normalize_python_package, ports::ProjectAnalyzer,
};

pub struct PythonProjectAnalyzer;

impl Default for PythonProjectAnalyzer {
    fn default() -> Self {
        Self::new()
    }
}

impl PythonProjectAnalyzer {
    pub fn new() -> Self {
        Self
    }
}

// PEP 621 format
#[derive(Deserialize, Default)]
struct ProjectSection {
    name: Option<String>,
    #[serde(default)]
    dependencies: Vec<String>,
}

// Poetry format
#[derive(Deserialize, Default)]
struct PoetrySection {
    name: Option<String>,
    #[serde(default)]
    dependencies: HashMap<String, toml::Value>,
}

#[derive(Deserialize, Default)]
struct ToolSection {
    #[serde(default)]
    poetry: PoetrySection,
}

#[derive(Deserialize)]
struct PyprojectToml {
    project: Option<ProjectSection>,
    #[serde(default)]
    tool: ToolSection,
}

fn extract_dep_name(dep: &str) -> &str {
    dep.split(&['>', '<', '=', '!', '[', ';', ' '][..])
        .next()
        .unwrap_or(dep)
        .trim()
}

fn normalize(name: &str) -> String {
    normalize_python_package(name)
}

impl ProjectAnalyzer for PythonProjectAnalyzer {
    fn analyze(&self, root: &Path) -> Result<CodeGraph, DomainError> {
        // 1. Scan immediate subdirectories for pyproject.toml
        let entries = std::fs::read_dir(root).map_err(|e| DomainError::IoError(e.to_string()))?;

        let mut packages: Vec<(String, String, Vec<String>)> = Vec::new(); // (dir, name, deps)

        for entry in entries.filter_map(|e| e.ok()) {
            let path = entry.path();
            if !path.is_dir() {
                continue;
            }
            let pyproject = path.join("pyproject.toml");
            if !pyproject.exists() {
                continue;
            }

            let content = std::fs::read_to_string(&pyproject)
                .map_err(|e| DomainError::IoError(e.to_string()))?;
            let parsed: PyprojectToml =
                toml::from_str(&content).map_err(|e| DomainError::ConfigError(e.to_string()))?;

            let dir_name = path
                .file_name()
                .unwrap_or_default()
                .to_string_lossy()
                .to_string();

            // Try PEP 621 [project] first, then Poetry [tool.poetry]
            let (name, deps) = if let Some(proj) = parsed.project {
                let name = proj.name.unwrap_or_else(|| dir_name.clone());
                let deps: Vec<String> = proj
                    .dependencies
                    .iter()
                    .map(|d| extract_dep_name(d).to_string())
                    .collect();
                (name, deps)
            } else if let Some(pname) = parsed.tool.poetry.name {
                let deps: Vec<String> = parsed
                    .tool
                    .poetry
                    .dependencies
                    .keys()
                    .filter(|k| k.as_str() != "python")
                    .cloned()
                    .collect();
                (pname, deps)
            } else {
                continue;
            };

            packages.push((dir_name, name, deps));
        }

        let known: HashSet<String> = packages
            .iter()
            .map(|(_, name, _)| normalize(name))
            .collect();

        let mut graph = CodeGraph::new();

        for (dir, name, _) in &packages {
            let file_path = FilePath::new(&format!("{}/pyproject.toml", dir))
                .map_err(|e| DomainError::IoError(e.to_string()))?;
            let element = CodeElement::new(name, CodeElementKind::Project, file_path, 1)?;
            graph.add_element(element);
        }

        for (_, pkg_name, deps) in &packages {
            for dep in deps {
                let dep_norm = normalize(dep);
                if known.contains(&dep_norm) && dep_norm != normalize(pkg_name) {
                    // find the canonical name (original casing) of the dep
                    if let Some((_, canonical, _)) =
                        packages.iter().find(|(_, n, _)| normalize(n) == dep_norm)
                        && let Ok(rel) =
                            Relationship::new(pkg_name, canonical, RelationshipKind::Composition)
                    {
                        graph.add_relationship(rel);
                    }
                }
            }
        }

        Ok(graph)
    }
}