mcp_methods/

list_dir.rs

//! Tree-formatted directory listing with optional per-entry annotation.
//!
//! Pure Rust. The Python wrapper in `mcp-methods-py` bridges a Python
//! callable into the [`ListDirOpts::annotate`] slot.

use ignore::overrides::OverrideBuilder;
use ignore::WalkBuilder;
use std::collections::{BTreeMap, HashMap, HashSet};
use std::path::{Path, PathBuf};
use std::sync::LazyLock;

const DEFAULT_SKIP_DIRS: &[&str] = &[
    ".git",
    "node_modules",
    "__pycache__",
    ".tox",
    ".mypy_cache",
    ".pytest_cache",
    "dist",
    "build",
    ".eggs",
    "venv",
    ".venv",
    "target",
    ".cargo",
    ".ruff_cache",
];

static DEFAULT_SKIP_SET: LazyLock<HashSet<&'static str>> =
    LazyLock::new(|| DEFAULT_SKIP_DIRS.iter().copied().collect());

/// An entry in the directory tree.
struct Entry {
    name: String,
    full_path: PathBuf,
    is_dir: bool,
    size: u64,
    children: BTreeMap<String, Entry>,
}

impl Entry {
    fn new_dir(name: String, full_path: PathBuf) -> Self {
        Self {
            name,
            full_path,
            is_dir: true,
            size: 0,
            children: BTreeMap::new(),
        }
    }

    fn new_file(name: String, full_path: PathBuf, size: u64) -> Self {
        Self {
            name,
            full_path,
            is_dir: false,
            size,
            children: BTreeMap::new(),
        }
    }
}

/// Optional knobs for [`list_dir`].
#[derive(Default)]
pub struct ListDirOpts<'a> {
    /// Recursion depth (default 1 = flat ls; 2+ = nested tree).
    pub depth: Option<usize>,
    pub glob: Option<&'a str>,
    pub dirs_only: bool,
    pub relative_to: Option<&'a str>,
    /// Respect `.gitignore` / `.git/info/exclude` rules (default `true`).
    pub respect_gitignore: bool,
    /// Custom directory names to skip. When `None`, a built-in list
    /// (`.git`, `node_modules`, `target`, …) is used.
    pub skip_dirs: Option<&'a [String]>,
    pub include_size: bool,
    /// Per-entry annotation callback. Receives the entry's path
    /// relative to `relative_to` (or the root if unset) and returns
    /// `Some(annotation)` to append after the entry, or `None` to skip.
    /// Used by the Python wrapper to bridge a Python callable.
    pub annotate: Option<&'a AnnotateFn>,
}

/// Closure signature for [`ListDirOpts::annotate`].
pub type AnnotateFn = dyn Fn(&str) -> Option<String>;

/// List directory contents with tree-formatted output.
pub fn list_dir(path: &str, opts: &ListDirOpts) -> Result<String, String> {
    let root = PathBuf::from(path)
        .canonicalize()
        .map_err(|e| format!("Cannot resolve '{}': {}", path, e))?;
    if !root.is_dir() {
        return Ok(format!("Error: '{}' is not a directory.", path));
    }

    let depth = opts.depth.unwrap_or(1);
    let respect_gitignore = opts.respect_gitignore;
    let glob = opts.glob;
    let relative_to = opts.relative_to;
    let dirs_only = opts.dirs_only;
    let include_size = opts.include_size;

    let custom_skip: Option<HashSet<String>> =
        opts.skip_dirs.map(|dirs| dirs.iter().cloned().collect());

    let mut tree = Entry::new_dir(dir_display_name(&root, relative_to), root.clone());
    let mut leaf_counts: BTreeMap<PathBuf, (usize, usize)> = BTreeMap::new();

    {
        let mut builder = WalkBuilder::new(&root);
        builder.max_depth(Some(depth + 1));
        builder.hidden(false);
        builder.git_ignore(respect_gitignore);
        builder.git_global(respect_gitignore);
        builder.git_exclude(respect_gitignore);

        if let Some(glob_pat) = glob {
            let mut overrides = OverrideBuilder::new(&root);
            overrides.add("*/").map_err(|e| format!("{}", e))?;
            overrides.add(glob_pat).map_err(|e| format!("{}", e))?;
            let built = overrides.build().map_err(|e| format!("{}", e))?;
            builder.overrides(built);
        }

        builder.filter_entry(move |entry| {
            if entry.file_type().map(|t| t.is_dir()).unwrap_or(false) {
                if let Some(name) = entry.file_name().to_str() {
                    return match &custom_skip {
                        Some(set) => !set.contains(name),
                        None => !DEFAULT_SKIP_SET.contains(name),
                    };
                }
            }
            true
        });

        for entry in builder.build().flatten() {
            let entry_path = entry.path().to_path_buf();
            if entry_path == root {
                continue;
            }
            let rel = match entry_path.strip_prefix(&root) {
                Ok(r) => r,
                Err(_) => continue,
            };
            let comp_count = rel.components().count();
            let is_dir = entry.file_type().map(|t| t.is_dir()).unwrap_or(false);

            if comp_count <= depth {
                // Build tree entry
                if dirs_only && !is_dir {
                    continue;
                }
                let components: Vec<String> = rel
                    .components()
                    .map(|c| c.as_os_str().to_string_lossy().to_string())
                    .collect();
                let size = if include_size && !is_dir {
                    entry.metadata().map(|m| m.len()).unwrap_or(0)
                } else {
                    0
                };
                insert_entry(&mut tree, &components, is_dir, size, &entry_path);
            } else {
                // depth+1: count for leaf directory summaries
                if let Some(parent) = entry_path.parent() {
                    let counter = leaf_counts.entry(parent.to_path_buf()).or_insert((0, 0));
                    if is_dir {
                        counter.0 += 1;
                    } else {
                        counter.1 += 1;
                    }
                }
            }
        }
    }

    // Prune empty dirs when glob is active
    if glob.is_some() && !dirs_only {
        prune_empty_dirs(&mut tree);
    }

    if tree.children.is_empty() {
        return Ok(format!("{}/ (empty)", tree.name));
    }

    // Collect annotations if callback provided
    let annotations = if let Some(annotate_fn) = opts.annotate {
        let mut map = HashMap::new();
        // Use relative_to base (canonicalized) so callback receives paths like
        // "networkx/algorithms/bridges.py" instead of just "bridges.py".
        // Falls back to root when relative_to is not set.
        let base = relative_to
            .and_then(|r| PathBuf::from(r).canonicalize().ok())
            .unwrap_or_else(|| root.clone());
        collect_annotations(&tree, &base, annotate_fn, &mut map);
        map
    } else {
        HashMap::new()
    };

    let mut output = Vec::new();
    output.push(format!("{}/", tree.name));
    render_tree(
        &tree,
        "",
        &mut output,
        include_size,
        &leaf_counts,
        &annotations,
    );
    Ok(output.join("\n"))
}

/// Recursively collect annotations for all entries in the tree.
fn collect_annotations(
    entry: &Entry,
    base: &Path,
    annotate_fn: &AnnotateFn,
    map: &mut HashMap<PathBuf, String>,
) {
    for child in entry.children.values() {
        let rel_path = child
            .full_path
            .strip_prefix(base)
            .unwrap_or(&child.full_path)
            .to_string_lossy()
            .to_string();
        if let Some(annotation) = annotate_fn(&rel_path) {
            map.insert(child.full_path.clone(), annotation);
        }
        if child.is_dir && !child.children.is_empty() {
            collect_annotations(child, base, annotate_fn, map);
        }
    }
}

fn insert_entry(
    tree: &mut Entry,
    components: &[String],
    is_dir: bool,
    size: u64,
    full_path: &Path,
) {
    let mut node = tree;
    for (i, comp) in components.iter().enumerate() {
        if i == components.len() - 1 {
            node.children.entry(comp.clone()).or_insert_with(|| {
                if is_dir {
                    Entry::new_dir(comp.clone(), full_path.to_path_buf())
                } else {
                    Entry::new_file(comp.clone(), full_path.to_path_buf(), size)
                }
            });
        } else {
            // Intermediate directory — construct its path
            let intermediate_path: PathBuf = full_path
                .components()
                .take(full_path.components().count() - (components.len() - 1 - i))
                .collect();
            node = node
                .children
                .entry(comp.clone())
                .or_insert_with(|| Entry::new_dir(comp.clone(), intermediate_path));
        }
    }
}

fn prune_empty_dirs(entry: &mut Entry) -> bool {
    if !entry.is_dir {
        return true;
    }
    entry.children.retain(|_, child| prune_empty_dirs(child));
    !entry.children.is_empty()
}

fn render_tree(
    entry: &Entry,
    prefix: &str,
    output: &mut Vec<String>,
    include_size: bool,
    leaf_counts: &BTreeMap<PathBuf, (usize, usize)>,
    annotations: &HashMap<PathBuf, String>,
) {
    let len = entry.children.len();

    // Pre-compute max name width for annotation alignment
    let max_name_width = if !annotations.is_empty() {
        entry
            .children
            .values()
            .map(|child| {
                let base = child.name.len() + if child.is_dir { 1 } else { 0 }; // +1 for "/"
                if include_size && !child.is_dir {
                    base + 2 + format_size(child.size).len() + 1 // "  (NNN B)"
                } else {
                    base
                }
            })
            .max()
            .unwrap_or(0)
    } else {
        0
    };

    for (i, child) in entry.children.values().enumerate() {
        let is_last = i == len - 1;
        let connector = if is_last { "└── " } else { "├── " };
        let child_prefix = if is_last { "    " } else { "│   " };

        if child.is_dir {
            let summary = if child.children.is_empty() {
                leaf_counts
                    .get(&child.full_path)
                    .map(|&(d, f)| format_summary(d, f))
                    .unwrap_or_default()
            } else {
                String::new()
            };
            let annotation = annotations.get(&child.full_path);
            if let Some(ann) = annotation {
                let name_part = format!("{}/", child.name);
                let pad = if max_name_width > name_part.len() {
                    max_name_width - name_part.len()
                } else {
                    0
                };
                output.push(format!(
                    "{}{}{}{}{}  {}",
                    prefix,
                    connector,
                    name_part,
                    summary,
                    " ".repeat(pad),
                    ann
                ));
            } else {
                output.push(format!("{}{}{}/{}", prefix, connector, child.name, summary));
            }
            if !child.children.is_empty() {
                render_tree(
                    child,
                    &format!("{}{}", prefix, child_prefix),
                    output,
                    include_size,
                    leaf_counts,
                    annotations,
                );
            }
        } else {
            let size_str = if include_size {
                format!("  ({})", format_size(child.size))
            } else {
                String::new()
            };
            let annotation = annotations.get(&child.full_path);
            if let Some(ann) = annotation {
                let name_part = format!("{}{}", child.name, size_str);
                let pad = if max_name_width > name_part.len() {
                    max_name_width - name_part.len()
                } else {
                    0
                };
                output.push(format!(
                    "{}{}{}{}  {}",
                    prefix,
                    connector,
                    name_part,
                    " ".repeat(pad),
                    ann
                ));
            } else {
                output.push(format!("{}{}{}{}", prefix, connector, child.name, size_str));
            }
        }
    }
}

fn format_summary(dirs: usize, files: usize) -> String {
    match (dirs, files) {
        (0, 0) => String::new(),
        (0, f) => format!("           [{} file{}]", f, if f == 1 { "" } else { "s" }),
        (d, 0) => format!("           [{} dir{}]", d, if d == 1 { "" } else { "s" }),
        (d, f) => format!(
            "           [{} dir{}, {} file{}]",
            d,
            if d == 1 { "" } else { "s" },
            f,
            if f == 1 { "" } else { "s" }
        ),
    }
}

fn format_size(bytes: u64) -> String {
    if bytes < 1024 {
        format!("{} B", bytes)
    } else if bytes < 1024 * 1024 {
        format!("{:.1} KB", bytes as f64 / 1024.0)
    } else if bytes < 1024 * 1024 * 1024 {
        format!("{:.1} MB", bytes as f64 / (1024.0 * 1024.0))
    } else {
        format!("{:.1} GB", bytes as f64 / (1024.0 * 1024.0 * 1024.0))
    }
}

fn dir_display_name(path: &Path, relative_to: Option<&str>) -> String {
    if let Some(base) = relative_to {
        let base_path = PathBuf::from(base);
        if let Ok(rel) = path.strip_prefix(&base_path) {
            let s = rel.to_string_lossy().to_string();
            if s.is_empty() {
                return path
                    .file_name()
                    .map(|n| n.to_string_lossy().to_string())
                    .unwrap_or_else(|| ".".to_string());
            }
            return s;
        }
    }
    path.file_name()
        .map(|n| n.to_string_lossy().to_string())
        .unwrap_or_else(|| ".".to_string())
}