feat: add support for enums with a rawvalue

elihwyma · Oct 30, 2024 · 009c212 · 009c212
1 parent ea619fc
commit 009c212
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 9 deletions.
diff --git a/Sources/PugiSwift/Decoding/XMLDecoderError.swift b/Sources/PugiSwift/Decoding/XMLDecoderError.swift
@@ -25,6 +25,8 @@ public enum XMLDecoderError: Error, LocalizedError, CustomStringConvertible {
 
     case failedToParseURL(string: String)
 
+    case invalidCase
+
     public var description: String {
         localizedDescription
     }
@@ -47,6 +49,8 @@ public enum XMLDecoderError: Error, LocalizedError, CustomStringConvertible {
             "Could not get name"
         case .failedToParseURL(string: let string):
             "Failed to parse into URL: \(string)"
+        case .invalidCase:
+            "Enum case not found"
         }
     }
 

diff --git a/Sources/PugiSwiftDemo/main.swift b/Sources/PugiSwiftDemo/main.swift
@@ -16,6 +16,7 @@ import PugiSwift
 @Node struct Record {
     let name: String
     let list: Int
+    let color: Colours
 }
 
 let str =
@@ -25,14 +26,24 @@ let str =
     <record>
         <name>Paul Koch</name>
         <list>17</list>
+        <color>red</color>
     </record>
     <record>
         <name>John Appleseed</name>
         <list>423</list>
+        <color>green</color>
     </record>
 </records> 
 """
 
+@Node enum Colours: String {
+
+    case red = "red"
+    case green = "green"
+    case blue = "blue"
+
+}
+
 do {
     let records = try Records(from: str)
     print(records)

diff --git a/Sources/PugiSwiftMacros/PugiSwiftMacros.swift b/Sources/PugiSwiftMacros/PugiSwiftMacros.swift
@@ -190,15 +190,7 @@ public struct NodeMacro: ExtensionMacro {
         }
     }
 
-    public static func expansion(of node: SwiftSyntax.AttributeSyntax,
-                                 attachedTo declaration: some SwiftSyntax.DeclGroupSyntax,
-                                 providingExtensionsOf type: some SwiftSyntax.TypeSyntaxProtocol,
-                                 conformingTo protocols: [SwiftSyntax.TypeSyntax],
-                                 in context: some SwiftSyntaxMacros.MacroExpansionContext) throws -> [SwiftSyntax.ExtensionDeclSyntax] {
-        // Validate its a struct we're working with
-        guard let structDecl = Struct(declaration) else {
-            throw MacroError("@Node can only be applied to structs")
-        }
+    private static func structExpansion(struct structDecl: Struct) throws -> [SwiftSyntax.ExtensionDeclSyntax] {
         // Get the access level of the struct
         let modifier = structDecl.accessLevel
         // Create the extension syntax
@@ -214,6 +206,45 @@ public struct NodeMacro: ExtensionMacro {
         return [extensionDecl]
     }
 
+    private static func enumExpansion(enum enumDecl: Enum) throws -> [SwiftSyntax.ExtensionDeclSyntax] {
+
+        var extensionDecl = try ExtensionDeclSyntax("extension \(raw: enumDecl.identifier): XMLDecodable {}")
+        guard let rawType = enumDecl.inheritedTypes.first?.normalizedDescription else {
+            throw MacroError("Enum must have a raw type.")
+        }
+
+        var functionDecl = Self.createFunction(with: enumDecl.accessLevel,
+                                               name: "node",
+                                               type: "PugiSwift.XMLNode")
+
+        let syntax = CodeBlockItemSyntax(
+        """
+        let rawValue = try \(raw: rawType).init(from: node)
+        guard let enumType = \(raw: enumDecl.identifier).init(rawValue: rawValue) else {
+            throw .invalidCase
+        }
+        self = enumType
+        """
+        )
+        functionDecl.body = CodeBlockSyntax(statements: CodeBlockItemListSyntax([syntax]))
+        let memberBlockItemSyntax = MemberBlockItemSyntax(decl: functionDecl)
+        extensionDecl.memberBlock.members.append(memberBlockItemSyntax)
+        return [extensionDecl]
+    }
+
+    public static func expansion(of node: SwiftSyntax.AttributeSyntax,
+                                 attachedTo declaration: some SwiftSyntax.DeclGroupSyntax,
+                                 providingExtensionsOf type: some SwiftSyntax.TypeSyntaxProtocol,
+                                 conformingTo protocols: [SwiftSyntax.TypeSyntax],
+                                 in context: some SwiftSyntaxMacros.MacroExpansionContext) throws -> [SwiftSyntax.ExtensionDeclSyntax] {
+        if let structDecl = Struct(declaration) {
+            return try Self.structExpansion(struct: structDecl)
+        } else if let enumDecl = Enum(declaration) {
+            return try Self.enumExpansion(enum: enumDecl)
+        }
+        throw MacroError("@Node applied to invalid type")
+    }
+
 }
 
 @main