Support Luau attributes

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Luau: Added support for parsing user-defined type functions ([#938](https://github.com/JohnnyMorganz/StyLua/issues/938))
+- Luau: Added support for parsing attributes (`@native` / `@deprecated`) on functions
 
 ### Fixed
 

src/formatters/functions.rs

@@ -8,7 +8,9 @@ use full_moon::ast::{
 use full_moon::tokenizer::{Token, TokenKind, TokenReference, TokenType};
 
 #[cfg(feature = "luau")]
-use crate::formatters::luau::{format_generic_declaration, format_type_specifier};
+use crate::formatters::luau::{
+    format_generic_declaration, format_luau_attribute, format_type_specifier,
+};
 use crate::{
     context::{
         create_function_call_trivia, create_function_definition_trivia, create_indent_trivia,
@@ -46,6 +48,18 @@ pub fn format_anonymous_function(
     shape: Shape,
 ) -> Box<AnonymousFunction> {
     const FUNCTION_LEN: usize = "function".len();
+
+    // Format attributes on one line, space separated
+    #[cfg(feature = "luau")]
+    let attributes = anonymous_function
+        .attributes()
+        .map(|attribute| {
+            format_luau_attribute(ctx, attribute, shape).update_trailing_trivia(
+                FormatTriviaType::Append(vec![Token::new(TokenType::spaces(1))]),
+            )
+        })
+        .collect();
+
     let function_definition_trivia = vec![create_function_definition_trivia(ctx)];
     let function_token = fmt_symbol!(ctx, anonymous_function.function_token(), "function", shape)
         .update_trailing_trivia(FormatTriviaType::Append(function_definition_trivia));
@@ -55,12 +69,15 @@ pub fn format_anonymous_function(
         shape.add_width(FUNCTION_LEN),
     );
 
-    Box::new(
-        anonymous_function
-            .clone()
-            .with_function_token(function_token)
-            .with_body(function_body),
-    )
+    let anonymous_function = anonymous_function
+        .clone()
+        .with_function_token(function_token)
+        .with_body(function_body);
+
+    #[cfg(feature = "luau")]
+    let anonymous_function = anonymous_function.with_attributes(attributes);
+
+    Box::new(anonymous_function)
 }
 
 /// An enum providing information regarding the next AST node after a function call.
@@ -1179,6 +1196,16 @@ pub fn format_function_declaration(
     let trailing_trivia = vec![create_newline_trivia(ctx)];
     let function_definition_trivia = vec![create_function_definition_trivia(ctx)];
 
+    #[cfg(feature = "luau")]
+    let attributes = function_declaration
+        .attributes()
+        .map(|attribute| {
+            format_luau_attribute(ctx, attribute, shape)
+                .update_leading_trivia(FormatTriviaType::Append(leading_trivia.clone()))
+                .update_trailing_trivia(FormatTriviaType::Append(trailing_trivia.clone()))
+        })
+        .collect();
+
     let function_token = fmt_symbol!(
         ctx,
         function_declaration.function_token(),
@@ -1193,9 +1220,14 @@ pub fn format_function_declaration(
     let function_body = format_function_body(ctx, function_declaration.body(), shape)
         .update_trailing_trivia(FormatTriviaType::Append(trailing_trivia));
 
-    FunctionDeclaration::new(formatted_function_name)
+    let function_declaration = FunctionDeclaration::new(formatted_function_name)
         .with_function_token(function_token)
-        .with_body(function_body)
+        .with_body(function_body);
+
+    #[cfg(feature = "luau")]
+    let function_declaration = function_declaration.with_attributes(attributes);
+
+    function_declaration
 }
 
 /// Formats a LocalFunction node
@@ -1209,6 +1241,15 @@ pub fn format_local_function(
     let trailing_trivia = vec![create_newline_trivia(ctx)];
     let function_definition_trivia = vec![create_function_definition_trivia(ctx)];
 
+    #[cfg(feature = "luau")]
+    let attributes = local_function
+        .attributes()
+        .map(|attribute| {
+            format_luau_attribute(ctx, attribute, shape)
+                .update_leading_trivia(FormatTriviaType::Append(leading_trivia.clone()))
+                .update_trailing_trivia(FormatTriviaType::Append(trailing_trivia.clone()))
+        })
+        .collect();
     let local_token = fmt_symbol!(ctx, local_function.local_token(), "local ", shape)
         .update_leading_trivia(FormatTriviaType::Append(leading_trivia));
     let function_token = fmt_symbol!(ctx, local_function.function_token(), "function ", shape);
@@ -1219,10 +1260,15 @@ pub fn format_local_function(
     let function_body = format_function_body(ctx, local_function.body(), shape)
         .update_trailing_trivia(FormatTriviaType::Append(trailing_trivia));
 
-    LocalFunction::new(formatted_name)
+    let local_function = LocalFunction::new(formatted_name)
         .with_local_token(local_token)
         .with_function_token(function_token)
-        .with_body(function_body)
+        .with_body(function_body);
+
+    #[cfg(feature = "luau")]
+    let local_function = local_function.with_attributes(attributes);
+
+    local_function
 }
 
 /// Formats a MethodCall node

src/formatters/luau.rs

@@ -27,9 +27,9 @@ use crate::{
 use full_moon::ast::{
     luau::{
         ExportedTypeDeclaration, ExportedTypeFunction, GenericDeclaration,
-        GenericDeclarationParameter, GenericParameterInfo, IndexedTypeInfo, TypeArgument,
-        TypeAssertion, TypeDeclaration, TypeField, TypeFieldKey, TypeFunction, TypeInfo,
-        TypeIntersection, TypeSpecifier, TypeUnion,
+        GenericDeclarationParameter, GenericParameterInfo, IndexedTypeInfo, LuauAttribute,
+        TypeArgument, TypeAssertion, TypeDeclaration, TypeField, TypeFieldKey, TypeFunction,
+        TypeInfo, TypeIntersection, TypeSpecifier, TypeUnion,
     },
     punctuated::Pair,
     CompoundAssignment, CompoundOp,
@@ -1574,3 +1574,14 @@ pub fn format_exported_type_function(
         .with_export_token(export_token)
         .with_type_function(type_function)
 }
+
+pub fn format_luau_attribute(
+    ctx: &Context,
+    attribute: &LuauAttribute,
+    shape: Shape,
+) -> LuauAttribute {
+    let at_sign = fmt_symbol!(ctx, attribute.at_sign(), "@", shape);
+    let name = format_token_reference(ctx, attribute.name(), shape);
+
+    attribute.clone().with_at_sign(at_sign).with_name(name)
+}

src/formatters/trivia.rs

@@ -3,8 +3,8 @@ use full_moon::ast::lua54::Attribute;
 #[cfg(feature = "luau")]
 use full_moon::ast::luau::{
     ElseIfExpression, GenericDeclaration, GenericDeclarationParameter, GenericParameterInfo,
-    IfExpression, IndexedTypeInfo, InterpolatedString, InterpolatedStringSegment, TypeArgument,
-    TypeAssertion, TypeDeclaration, TypeField, TypeFieldKey, TypeFunction, TypeInfo,
+    IfExpression, IndexedTypeInfo, InterpolatedString, InterpolatedStringSegment, LuauAttribute,
+    TypeArgument, TypeAssertion, TypeDeclaration, TypeField, TypeFieldKey, TypeFunction, TypeInfo,
     TypeIntersection, TypeSpecifier, TypeUnion,
 };
 use full_moon::ast::{
@@ -1114,3 +1114,10 @@ define_update_leading_trivia!(InterpolatedStringSegment, |this, leading| {
         expression: this.expression.to_owned(),
     }
 });
+
+#[cfg(feature = "luau")]
+define_update_trivia!(LuauAttribute, |this, leading, trailing| {
+    this.clone()
+        .with_at_sign(this.at_sign().update_leading_trivia(leading))
+        .with_name(this.name().update_trailing_trivia(trailing))
+});

test/example/test.luau

@@ -0,0 +1,13 @@
+local types = require("shared/types")
+
+local function Area() end
+
+local z = [==[testing]==]
+
+return {
+	Area {
+		act = {},
+		name = "",
+		description = "",
+	},
+}

tests/inputs-luau/attributes-1.lua

@@ -0,0 +1,10 @@
+@native
+function foo()
+end
+
+@deprecated
+local function bar()
+end
+
+local x = @native function()
+end

tests/inputs-luau/attributes-2.lua

@@ -0,0 +1,7 @@
+	@native
+function foo()
+end
+
+@native @deprecated
+function bar()
+end

tests/inputs-luau/attributes-3.lua

@@ -0,0 +1,5 @@
+local x =    @native  function()
+end
+
+local y = @native   @deprecated function()
+end

tests/snapshots/tests__luau@attributes-1.lua.snap

@@ -0,0 +1,13 @@
+---
+source: tests/tests.rs
+expression: "format(&contents, LuaVersion::Luau)"
+input_file: tests/inputs-luau/attributes-1.lua
+snapshot_kind: text
+---
+@native
+function foo() end
+
+@deprecated
+local function bar() end
+
+local x = @native function() end

tests/snapshots/tests__luau@attributes-2.lua.snap

@@ -0,0 +1,12 @@
+---
+source: tests/tests.rs
+expression: "format(&contents, LuaVersion::Luau)"
+input_file: tests/inputs-luau/attributes-2.lua
+snapshot_kind: text
+---
+@native
+function foo() end
+
+@native
+@deprecated
+function bar() end

tests/snapshots/tests__luau@attributes-3.lua.snap

@@ -0,0 +1,9 @@
+---
+source: tests/tests.rs
+expression: "format(&contents, LuaVersion::Luau)"
+input_file: tests/inputs-luau/attributes-3.lua
+snapshot_kind: text
+---
+local x = @native function() end
+
+local y = @native @deprecated function() end