rustledger-lsp 0.15.0

Language Server Protocol implementation for Beancount
Documentation
//! Document color handler for visual amount feedback.
//!
//! Provides color information for:
//! - Negative amounts: red
//! - Positive amounts: green
//! - Zero amounts: gray

use lsp_types::{
    Color, ColorInformation, ColorPresentation, ColorPresentationParams, DocumentColorParams,
    Position, Range,
};
use rustledger_core::Directive;
use rustledger_parser::ParseResult;

use super::utils::byte_offset_to_position;

/// Red color for negative amounts.
const COLOR_NEGATIVE: Color = Color {
    red: 0.9,
    green: 0.2,
    blue: 0.2,
    alpha: 1.0,
};

/// Green color for positive amounts.
const COLOR_POSITIVE: Color = Color {
    red: 0.2,
    green: 0.8,
    blue: 0.3,
    alpha: 1.0,
};

/// Gray color for zero amounts.
const COLOR_ZERO: Color = Color {
    red: 0.5,
    green: 0.5,
    blue: 0.5,
    alpha: 1.0,
};

/// Handle a document color request.
pub fn handle_document_color(
    _params: &DocumentColorParams,
    source: &str,
    parse_result: &ParseResult,
) -> Option<Vec<ColorInformation>> {
    let mut colors = Vec::new();

    for spanned in &parse_result.directives {
        match &spanned.value {
            Directive::Transaction(txn) => {
                let (start_line, _) = byte_offset_to_position(source, spanned.span.start);

                for (i, posting) in txn.postings.iter().enumerate() {
                    if let Some(units) = &posting.units
                        && let Some(number) = units.number()
                    {
                        let posting_line = start_line + 1 + i as u32;
                        let line_text = source.lines().nth(posting_line as usize).unwrap_or("");

                        // Find the amount in the line
                        let amount_str = number.to_string();
                        if let Some(range) = find_amount_range(line_text, &amount_str, posting_line)
                        {
                            let color = if number.is_sign_negative() {
                                COLOR_NEGATIVE
                            } else if number.is_zero() {
                                COLOR_ZERO
                            } else {
                                COLOR_POSITIVE
                            };

                            colors.push(ColorInformation { range, color });
                        }
                    }
                }
            }
            Directive::Balance(bal) => {
                let (line, _) = byte_offset_to_position(source, spanned.span.start);
                let line_text = source.lines().nth(line as usize).unwrap_or("");

                let amount_str = bal.amount.number.to_string();
                if let Some(range) = find_amount_range(line_text, &amount_str, line) {
                    let color = if bal.amount.number.is_sign_negative() {
                        COLOR_NEGATIVE
                    } else if bal.amount.number.is_zero() {
                        COLOR_ZERO
                    } else {
                        COLOR_POSITIVE
                    };

                    colors.push(ColorInformation { range, color });
                }
            }
            Directive::Price(price) => {
                let (line, _) = byte_offset_to_position(source, spanned.span.start);
                let line_text = source.lines().nth(line as usize).unwrap_or("");

                let amount_str = price.amount.number.to_string();
                if let Some(range) = find_amount_range(line_text, &amount_str, line) {
                    colors.push(ColorInformation {
                        range,
                        color: COLOR_POSITIVE, // Prices are always "positive" in context
                    });
                }
            }
            _ => {}
        }
    }

    if colors.is_empty() {
        None
    } else {
        Some(colors)
    }
}

/// Handle a color presentation request.
/// This is called when the user wants to change a color (not really applicable for amounts).
pub fn handle_color_presentation(params: &ColorPresentationParams) -> Vec<ColorPresentation> {
    // We don't support changing colors - amounts are data, not colors
    // Just return the current representation
    let label = if params.color.red > 0.5 && params.color.green < 0.5 {
        "Negative amount"
    } else if params.color.green > 0.5 {
        "Positive amount"
    } else {
        "Zero amount"
    };

    vec![ColorPresentation {
        label: label.to_string(),
        text_edit: None,
        additional_text_edits: None,
    }]
}

/// Find the range of an amount in a line.
fn find_amount_range(line: &str, amount_str: &str, line_num: u32) -> Option<Range> {
    // Look for the amount pattern (may have negative sign)
    let search_patterns = [
        amount_str.to_string(),
        format!("-{}", amount_str.trim_start_matches('-')),
    ];

    for pattern in &search_patterns {
        if let Some(pos) = line.find(pattern) {
            // Verify it's a standalone number (not part of a larger string)
            let before_ok = pos == 0
                || !line
                    .chars()
                    .nth(pos - 1)
                    .unwrap_or(' ')
                    .is_ascii_alphanumeric();
            let after_pos = pos + pattern.len();
            let after_ok = after_pos >= line.len()
                || !line.chars().nth(after_pos).unwrap_or(' ').is_ascii_digit();

            if before_ok && after_ok {
                return Some(Range {
                    start: Position::new(line_num, pos as u32),
                    end: Position::new(line_num, (pos + pattern.len()) as u32),
                });
            }
        }
    }

    None
}

#[cfg(test)]
mod tests {
    use super::*;
    use rustledger_parser::parse;

    #[test]
    fn test_document_color_positive_negative() {
        let source = r#"2024-01-15 * "Coffee"
  Assets:Bank  -5.00 USD
  Expenses:Food  5.00 USD
"#;
        let result = parse(source);
        let params = DocumentColorParams {
            text_document: lsp_types::TextDocumentIdentifier {
                uri: "file:///test.beancount".parse().unwrap(),
            },
            work_done_progress_params: Default::default(),
            partial_result_params: Default::default(),
        };

        let colors = handle_document_color(&params, source, &result);
        assert!(colors.is_some());

        let colors = colors.unwrap();
        assert_eq!(colors.len(), 2);

        // First posting is negative (red)
        assert!(colors[0].color.red > 0.5);
        assert!(colors[0].color.green < 0.5);

        // Second posting is positive (green)
        assert!(colors[1].color.green > 0.5);
        assert!(colors[1].color.red < 0.5);
    }

    #[test]
    fn test_document_color_balance() {
        let source = r#"2024-01-31 balance Assets:Bank 100 USD
"#;
        let result = parse(source);
        let params = DocumentColorParams {
            text_document: lsp_types::TextDocumentIdentifier {
                uri: "file:///test.beancount".parse().unwrap(),
            },
            work_done_progress_params: Default::default(),
            partial_result_params: Default::default(),
        };

        let colors = handle_document_color(&params, source, &result);
        assert!(colors.is_some());

        let colors = colors.unwrap();
        assert_eq!(colors.len(), 1);
        // Positive balance (green)
        assert!(colors[0].color.green > 0.5);
    }
}