ruff/crates/ruff_python_formatter/src/context.rs

use crate::comments::Comments;
use ruff_formatter::{FormatContext, SimpleFormatOptions, SourceCode};
use ruff_python_ast::source_code::Locator;
use std::fmt::{Debug, Formatter};

#[derive(Clone)]
pub struct PyFormatContext<'a> {
    options: SimpleFormatOptions,
    contents: &'a str,
    comments: Comments<'a>,
    node_level: NodeLevel,
}

impl<'a> PyFormatContext<'a> {
    pub(crate) fn new(
        options: SimpleFormatOptions,
        contents: &'a str,
        comments: Comments<'a>,
    ) -> Self {
        Self {
            options,
            contents,
            comments,
            node_level: NodeLevel::TopLevel,
        }
    }

    pub(crate) fn contents(&self) -> &'a str {
        self.contents
    }

    #[allow(unused)]
    pub(crate) fn locator(&self) -> Locator<'a> {
        Locator::new(self.contents)
    }

    pub(crate) fn set_node_level(&mut self, level: NodeLevel) {
        self.node_level = level;
    }

    pub(crate) fn node_level(&self) -> NodeLevel {
        self.node_level
    }

    pub(crate) fn comments(&self) -> &Comments<'a> {
        &self.comments
    }
}

impl FormatContext for PyFormatContext<'_> {
    type Options = SimpleFormatOptions;

    fn options(&self) -> &Self::Options {
        &self.options
    }

    fn source_code(&self) -> SourceCode {
        SourceCode::new(self.contents)
    }
}

impl Debug for PyFormatContext<'_> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PyFormatContext")
            .field("options", &self.options)
            .field("comments", &self.comments.debug(self.source_code()))
            .field("node_level", &self.node_level)
            .field("source", &self.contents)
            .finish()
    }
}

/// What's the enclosing level of the outer node.
#[derive(Copy, Clone, Debug, Eq, PartialEq, Default)]
pub(crate) enum NodeLevel {
    /// Formatting statements on the module level.
    #[default]
    TopLevel,

    /// Formatting the body statements of a [compound statement](https://docs.python.org/3/reference/compound_stmts.html#compound-statements)
    /// (`if`, `while`, `match`, etc.).
    CompoundStatement,

    /// Formatting nodes that are enclosed in a parenthesized expression.
    Expression,
}