Handle 5.4 attributes

src/formatters/assignment.rs

@@ -9,6 +9,8 @@ use full_moon::{
     tokenizer::TokenType,
 };
 
+#[cfg(feature = "lua54")]
+use crate::formatters::lua54::format_attribute;
 #[cfg(feature = "luau")]
 use crate::formatters::luau::format_type_specifier;
 use crate::{
@@ -386,13 +388,21 @@ fn format_local_no_assignment(
         Some(1),
     );
 
+    #[cfg(feature = "lua54")]
+    let attributes = assignment
+        .attributes()
+        .map(|x| x.map(|attribute| format_attribute(ctx, attribute, shape)))
+        .collect();
+
     #[cfg(feature = "luau")]
     let type_specifiers: Vec<Option<TypeSpecifier>> = assignment
         .type_specifiers()
         .map(|x| x.map(|type_specifier| format_type_specifier(ctx, type_specifier, shape)))
         .collect();
 
     let local_assignment = LocalAssignment::new(name_list);
+    #[cfg(feature = "lua54")]
+    let local_assignment = local_assignment.with_attributes(attributes);
     #[cfg(feature = "luau")]
     let local_assignment = local_assignment.with_type_specifiers(type_specifiers);
 
@@ -439,21 +449,30 @@ pub fn format_local_assignment_no_trivia(
             format_expression,
         );
 
+        #[cfg(feature = "lua54")]
+        let attributes: Vec<Option<_>> = assignment
+            .attributes()
+            .map(|x| x.map(|attribute| format_attribute(ctx, attribute, shape)))
+            .collect();
+
         #[cfg(feature = "luau")]
         let type_specifiers: Vec<Option<TypeSpecifier>> = assignment
             .type_specifiers()
             .map(|x| x.map(|type_specifier| format_type_specifier(ctx, type_specifier, shape)))
             .collect();
-        let type_specifier_len;
-        #[cfg(feature = "luau")]
+
+        let mut type_specifier_len = 0;
+        #[cfg(feature = "lua54")]
         {
-            type_specifier_len = type_specifiers.iter().fold(0, |acc, x| {
+            type_specifier_len += attributes.iter().fold(0, |acc, x| {
                 acc + x.as_ref().map_or(0, |y| y.to_string().len())
             });
         }
-        #[cfg(not(feature = "luau"))]
+        #[cfg(feature = "luau")]
         {
-            type_specifier_len = 0;
+            type_specifier_len += type_specifiers.iter().fold(0, |acc, x| {
+                acc + x.as_ref().map_or(0, |y| y.to_string().len())
+            });
         }
 
         // Test the assignment to see if its over width
@@ -483,6 +502,8 @@ pub fn format_local_assignment_no_trivia(
         }
 
         let local_assignment = LocalAssignment::new(name_list);
+        #[cfg(feature = "lua54")]
+        let local_assignment = local_assignment.with_attributes(attributes);
         #[cfg(feature = "luau")]
         let local_assignment = local_assignment.with_type_specifiers(type_specifiers);
         local_assignment

src/formatters/lua54.rs

@@ -0,0 +1,21 @@
+use crate::{
+    context::Context,
+    formatters::{
+        general::{format_contained_span, format_token_reference},
+        trivia::{FormatTriviaType, UpdateLeadingTrivia},
+    },
+    shape::Shape,
+};
+use full_moon::{
+    ast::lua54::Attribute,
+    tokenizer::{Token, TokenType},
+};
+
+pub fn format_attribute(ctx: &Context, attribute: &Attribute, shape: Shape) -> Attribute {
+    let brackets = format_contained_span(ctx, attribute.brackets(), shape).update_leading_trivia(
+        FormatTriviaType::Append(vec![Token::new(TokenType::spaces(1))]),
+    );
+    let name = format_token_reference(ctx, attribute.name(), shape);
+
+    Attribute::new(name).with_brackets(brackets)
+}

src/formatters/mod.rs

@@ -9,6 +9,8 @@ pub mod expression;
 pub mod functions;
 #[cfg(feature = "lua52")]
 pub mod lua52;
+#[cfg(feature = "lua54")]
+pub mod lua54;
 #[cfg(feature = "luau")]
 pub mod luau;
 pub mod stmt;

src/formatters/trivia.rs

@@ -1,3 +1,5 @@
+#[cfg(feature = "lua54")]
+use full_moon::ast::lua54::Attribute;
 #[cfg(feature = "luau")]
 use full_moon::ast::types::{
     ElseIfExpression, GenericDeclaration, GenericDeclarationParameter, GenericParameterInfo,
@@ -490,6 +492,12 @@ define_update_trivia!(Assignment, |this, leading, trailing| {
         .with_expressions(this.expressions().update_trailing_trivia(trailing))
 });
 
+#[cfg(feature = "lua54")]
+define_update_trivia!(Attribute, |this, leading, trailing| {
+    this.to_owned()
+        .with_brackets(this.brackets().update_trivia(leading, trailing))
+});
+
 define_update_trivia!(LocalAssignment, |this, leading, trailing| {
     if this.expressions().is_empty() {
         // Handle if the last item had a type specifier set
@@ -507,6 +515,20 @@ define_update_trivia!(LocalAssignment, |this, leading, trailing| {
             }
         );
 
+        cfg_if::cfg_if!(
+            if #[cfg(feature = "lua54")] {
+                let mut attributes = this.attributes().map(|x| x.cloned()).collect::<Vec<_>>();
+
+                if let Some(Some(attribute)) = attributes.pop() {
+                    attributes.push(Some(attribute.update_trailing_trivia(trailing)));
+
+                    return this.clone()
+                        .with_local_token(this.local_token().update_leading_trivia(leading))
+                        .with_attributes(attributes);
+                }
+            }
+        );
+
         this.clone()
             .with_local_token(this.local_token().update_leading_trivia(leading))
             .with_names(this.names().update_trailing_trivia(trailing))

src/formatters/trivia_util.rs

@@ -8,12 +8,12 @@ use full_moon::ast::span::ContainedSpan;
 #[cfg(feature = "luau")]
 use full_moon::ast::types::{
     GenericDeclarationParameter, GenericParameterInfo, IndexedTypeInfo, TypeArgument,
-    TypeDeclaration, TypeInfo,
+    TypeDeclaration, TypeInfo, TypeSpecifier,
 };
 use full_moon::{
     ast::{
-        BinOp, Block, Call, Expression, Field, FunctionArgs, Index, LastStmt, Parameter, Prefix,
-        Stmt, Suffix, TableConstructor, UnOp, Value, Var,
+        BinOp, Block, Call, Expression, Field, FunctionArgs, Index, LastStmt, LocalAssignment,
+        Parameter, Prefix, Stmt, Suffix, TableConstructor, UnOp, Value, Var,
     },
     node::Node,
     tokenizer::{Token, TokenKind, TokenReference, TokenType},
@@ -884,6 +884,81 @@ fn get_type_declaration_trailing_trivia(
     )
 }
 
+#[cfg(feature = "luau")]
+fn type_specifier_trailing_trivia(type_specifier: &TypeSpecifier) -> Vec<Token> {
+    get_type_info_trailing_trivia(type_specifier.type_info().clone()).1
+}
+
+fn get_empty_local_assignment_trailing_trivia(
+    local_assignment: LocalAssignment,
+) -> (LocalAssignment, Vec<Token>) {
+    let mut trailing_trivia = Vec::new();
+
+    #[cfg(feature = "luau")]
+    {
+        let mut type_specifiers = local_assignment
+            .type_specifiers()
+            .map(|x| x.cloned())
+            .collect::<Vec<_>>();
+
+        if let Some(Some(type_specifier)) = type_specifiers.pop() {
+            trailing_trivia = type_specifier_trailing_trivia(&type_specifier);
+
+            type_specifiers.push(Some(
+                type_specifier.update_trailing_trivia(FormatTriviaType::Replace(vec![])),
+            ));
+
+            return (
+                local_assignment.with_type_specifiers(type_specifiers),
+                trailing_trivia,
+            );
+        }
+    }
+
+    #[cfg(feature = "lua54")]
+    {
+        let mut attributes = local_assignment
+            .attributes()
+            .map(|x| x.cloned())
+            .collect::<Vec<_>>();
+
+        if let Some(Some(attribute)) = attributes.pop() {
+            trailing_trivia = attribute
+                .brackets()
+                .tokens()
+                .1
+                .trailing_trivia()
+                .cloned()
+                .collect();
+
+            attributes.push(Some(
+                attribute.update_trailing_trivia(FormatTriviaType::Replace(vec![])),
+            ));
+
+            return (
+                local_assignment.with_attributes(attributes),
+                trailing_trivia,
+            );
+        }
+    }
+
+    // Unassigned local variable
+    let mut formatted_name_list = local_assignment.names().to_owned();
+    // Retrieve last item and take its trailing comments
+    if let Some(last_pair) = formatted_name_list.pop() {
+        let pair = last_pair.map(|value| {
+            trailing_trivia = value.trailing_trivia().map(|x| x.to_owned()).collect();
+            value.update_trailing_trivia(FormatTriviaType::Replace(vec![]))
+        });
+        formatted_name_list.push(pair);
+    }
+
+    (
+        local_assignment.with_names(formatted_name_list),
+        trailing_trivia,
+    )
+}
+
 pub fn get_stmt_trailing_trivia(stmt: Stmt) -> (Stmt, Vec<Token>) {
     let (updated_stmt, trailing_trivia) = match stmt {
         Stmt::Assignment(assignment) => {
@@ -906,17 +981,10 @@ pub fn get_stmt_trailing_trivia(stmt: Stmt) -> (Stmt, Vec<Token>) {
         Stmt::LocalAssignment(local_assignment) => {
             let mut trailing_trivia = Vec::new();
             let new_assignment = if local_assignment.expressions().is_empty() {
-                // Unassigned local variable
-                let mut formatted_name_list = local_assignment.names().to_owned();
-                // Retrieve last item and take its trailing comments
-                if let Some(last_pair) = formatted_name_list.pop() {
-                    let pair = last_pair.map(|value| {
-                        trailing_trivia = value.trailing_trivia().map(|x| x.to_owned()).collect();
-                        value.update_trailing_trivia(FormatTriviaType::Replace(vec![]))
-                    });
-                    formatted_name_list.push(pair);
-                }
-                local_assignment.with_names(formatted_name_list)
+                let (assignment, trivia) =
+                    get_empty_local_assignment_trailing_trivia(local_assignment);
+                trailing_trivia = trivia;
+                assignment
             } else {
                 // Add newline at the end of LocalAssignment expression list
                 // Expression list should already be formatted

tests/tests.rs

@@ -26,6 +26,7 @@ fn test_full_moon_test_suite() {
 #[cfg(feature = "luau")]
 fn test_luau() {
     insta::glob!("inputs-luau/*.lua", |path| {
+        dbg!(path);
         let contents = std::fs::read_to_string(path).unwrap();
         insta::assert_snapshot!(format(&contents));
     })