stylua 0.4.0

A code formatter for Lua
Documentation
use full_moon::ast::{punctuated::Punctuated, Assignment, LocalAssignment};
use full_moon::tokenizer::TokenReference;
use std::borrow::Cow;

use crate::formatters::{
    trivia_formatter::{self, FormatTriviaType},
    CodeFormatter,
};

impl CodeFormatter {
    pub fn format_assignment<'ast>(&mut self, assignment: &Assignment<'ast>) -> Assignment<'ast> {
        let (var_list, mut var_comments_buf) =
            self.format_punctuated(assignment.var_list(), &CodeFormatter::format_var);

        let (mut expr_list, mut expr_comments_buf) =
            self.format_punctuated(assignment.expr_list(), &CodeFormatter::format_expression);

        match expr_list.pop() {
            Some(pair) => {
                var_comments_buf.append(&mut expr_comments_buf);

                let pair = pair.map(|expr| {
                    trivia_formatter::expression_add_trailing_trivia(
                        expr,
                        FormatTriviaType::Append(var_comments_buf),
                    )
                });
                expr_list.push(pair);
            }
            None => panic!("assignment with no expression"),
        }

        Assignment::new(var_list, expr_list)
            .with_equal_token(Cow::Owned(TokenReference::symbol(" = ").unwrap()))
    }

    pub fn format_local_assignment<'ast>(
        &mut self,
        assignment: &LocalAssignment<'ast>,
    ) -> LocalAssignment<'ast> {
        let local_token = crate::fmt_symbol!(self, assignment.local_token(), "local ");
        let (mut name_list, mut name_list_comments_buf) = self.format_punctuated(
            assignment.name_list(),
            &CodeFormatter::format_token_reference_mut,
        );

        #[cfg(feature = "luau")]
        let type_specifiers = assignment
            .type_specifiers()
            .map(|x| match x {
                Some(type_specifier) => Some(self.format_type_specifier(type_specifier)),
                None => None,
            })
            .collect();

        if assignment.expr_list().is_empty() {
            if let Some(pair) = name_list.pop() {
                let pair = pair.map(|name| {
                    Cow::Owned(trivia_formatter::token_reference_add_trivia(
                        name.to_owned().into_owned(),
                        FormatTriviaType::NoChange,
                        FormatTriviaType::Append(name_list_comments_buf),
                    ))
                });
                name_list.push(pair);
            }

            let local_assignment = LocalAssignment::new(name_list)
                .with_local_token(local_token)
                .with_equal_token(None)
                .with_expr_list(Punctuated::new());

            #[cfg(feature = "luau")]
            let local_assignment = local_assignment.with_type_specifiers(type_specifiers);
            local_assignment
        } else {
            let equal_token = crate::fmt_symbol!(self, assignment.equal_token().unwrap(), " = ");
            let (mut expr_list, mut expr_comments_buf) =
                self.format_punctuated(assignment.expr_list(), &CodeFormatter::format_expression);

            if let Some(pair) = expr_list.pop() {
                name_list_comments_buf.append(&mut expr_comments_buf);

                let pair = pair.map(|expr| {
                    trivia_formatter::expression_add_trailing_trivia(
                        expr,
                        FormatTriviaType::Append(name_list_comments_buf),
                    )
                });
                expr_list.push(pair);
            }

            let local_assignment = LocalAssignment::new(name_list)
                .with_local_token(local_token)
                .with_equal_token(Some(equal_token))
                .with_expr_list(expr_list);
            #[cfg(feature = "luau")]
            let local_assignment = local_assignment.with_type_specifiers(type_specifiers);

            local_assignment
        }
    }
}