From 2b888d052c6c3de7bdfd2b37a3780b237996682c Mon Sep 17 00:00:00 2001
From: Nayanda Haberty <hainayanda@outlook.com>
Date: Sun, 17 Mar 2024 14:34:40 +0700
Subject: [PATCH] Add open class and subclass to control stub modifier output

---
 Sources/SwiftEnvironment/Macros.swift                  |  2 ++
 .../SwiftEnvironmentMacro/Model/StubDeclaration.swift  |  6 +++---
 Sources/SwiftEnvironmentMacro/StubGeneratorMacro.swift | 10 +++++++---
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/Sources/SwiftEnvironment/Macros.swift b/Sources/SwiftEnvironment/Macros.swift
index 3e3f009..e720d68 100644
--- a/Sources/SwiftEnvironment/Macros.swift
+++ b/Sources/SwiftEnvironment/Macros.swift
@@ -25,7 +25,9 @@ public macro Stubbed(name: String? = nil, type: StubType = .struct, _ values: De
 public enum StubType {
     case `struct`
     case `class`
+    case openClass
     case subclass(AnyClass.Type)
+    case openSubclass(AnyClass.Type)
 }
 
 public struct DefaultType {
diff --git a/Sources/SwiftEnvironmentMacro/Model/StubDeclaration.swift b/Sources/SwiftEnvironmentMacro/Model/StubDeclaration.swift
index b129e86..1f2de4b 100644
--- a/Sources/SwiftEnvironmentMacro/Model/StubDeclaration.swift
+++ b/Sources/SwiftEnvironmentMacro/Model/StubDeclaration.swift
@@ -29,13 +29,13 @@ struct StubDeclaration: CustomStringConvertible {
 extension StubDeclaration {
     
     enum InstanceType {
-        case `class`(`protocol`: String, superClass: String? = nil)
+        case `class`(`protocol`: String, superClass: String? = nil, final: Bool = true)
         case `struct`(`protocol`: String)
         
         func declarationString(for name: String) -> String {
             switch self {
-            case .class(let protocolName, let superClass):
-                return "final class \(name): \(superClassClause(for: superClass))\(protocolName)"
+            case .class(let protocolName, let superClass, let final):
+                return "\(final ? "final ": "")class \(name): \(superClassClause(for: superClass))\(protocolName)"
             case .struct(let protocolName):
                 return "struct \(name): \(protocolName)"
             }
diff --git a/Sources/SwiftEnvironmentMacro/StubGeneratorMacro.swift b/Sources/SwiftEnvironmentMacro/StubGeneratorMacro.swift
index fc3df58..8e6b2d4 100644
--- a/Sources/SwiftEnvironmentMacro/StubGeneratorMacro.swift
+++ b/Sources/SwiftEnvironmentMacro/StubGeneratorMacro.swift
@@ -108,21 +108,25 @@ private extension AttributeSyntax {
                 }
                 return .struct(protocol: protocolName)
             } else if argument.trimmedDescription.match(#"^(StubType)?\.`?class`?$"#) {
-                return .class(protocol: protocolName, superClass: nil)
+                return .class(protocol: protocolName, superClass: nil, final: true)
+            } else if argument.trimmedDescription.match(#"^(StubType)?\.openClass$"#) {
+                return .class(protocol: protocolName, superClass: nil, final: false)
             } else {
                 throw StubGeneratorMacroError.cannotDetermineDefaultValue(argument.trimmedDescription)
             }
         }
     
     func argsTypeArgument(for protocolName: String, argument: FunctionCallExprSyntax) throws -> StubDeclaration.InstanceType {
-        guard argument.calledExpression.trimmedDescription.match(#"^(StubType)?\.`?subclass`?$"#),
+        let isSubclass: Bool = argument.calledExpression.trimmedDescription.match(#"^(StubType)?\.subclass$"#)
+        let isOpenSubclass: Bool = argument.calledExpression.trimmedDescription.match(#"^(StubType)?\.openSubclass$"#)
+        guard isSubclass || isOpenSubclass,
               let argumentExpression = argument.arguments.first?.expression,
               argumentExpression.trimmedDescription.match(#"^\S+\.self$"#),
               let memberAccess = argumentExpression.as(MemberAccessExprSyntax.self),
               let superClass = memberAccess.base?.trimmedDescription else {
             throw StubGeneratorMacroError.cannotDetermineDefaultValue(argument.trimmedDescription)
         }
-        return .class(protocol: protocolName, superClass: superClass)
+        return .class(protocol: protocolName, superClass: superClass, final: isSubclass)
     }
 }