Allow @Unredacted classes & @Redacted child objects (#212)

ZacSweers · Apr 26, 2024 · 79efe4c · 79efe4c
1 parent 085a2e0
commit 79efe4c
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 6 deletions.
diff --git a/...plugin-annotations/src/commonMain/kotlin/dev/zacsweers/redacted/annotations/Unredacted.kt b/...plugin-annotations/src/commonMain/kotlin/dev/zacsweers/redacted/annotations/Unredacted.kt
@@ -16,6 +16,7 @@
 package dev.zacsweers.redacted.annotations
 
 import kotlin.annotation.AnnotationRetention.BINARY
+import kotlin.annotation.AnnotationTarget.CLASS
 import kotlin.annotation.AnnotationTarget.PROPERTY
 
 /**
@@ -33,4 +34,4 @@ import kotlin.annotation.AnnotationTarget.PROPERTY
  * println(user) // User(name = "Bob", phoneNumber = "██")
  * ```
  */
-@Retention(BINARY) @Target(PROPERTY) public annotation class Unredacted
+@Retention(BINARY) @Target(PROPERTY, CLASS) public annotation class Unredacted
diff --git a/...cted-compiler-plugin/src/main/kotlin/dev/zacsweers/redacted/compiler/RedactedIrVisitor.kt b/...cted-compiler-plugin/src/main/kotlin/dev/zacsweers/redacted/compiler/RedactedIrVisitor.kt
@@ -87,6 +87,7 @@ internal class RedactedIrVisitor(
 
       val properties = mutableListOf<Property>()
       val classIsRedacted = declarationParent.hasAnnotation(redactedAnnotation)
+      val classIsUnredacted = declarationParent.hasAnnotation(unredactedAnnotation)
       val supertypeIsRedacted by
         lazy(NONE) {
           declarationParent.getAllSuperclasses().any { it.hasAnnotation(redactedAnnotation) }
@@ -106,7 +107,7 @@ internal class RedactedIrVisitor(
         properties += Property(prop, isRedacted, isUnredacted, parameter)
       }
 
-      if (classIsRedacted || supertypeIsRedacted || anyRedacted) {
+      if (classIsRedacted || supertypeIsRedacted || classIsUnredacted || anyRedacted) {
         if (declaration.origin == IrDeclarationOrigin.DEFINED) {
           declaration.reportError(
             "@Redacted is only supported on data or value classes that do *not* have a custom toString() function. Please remove the function or remove the @Redacted annotations."
@@ -134,7 +135,24 @@ internal class RedactedIrVisitor(
           return super.visitFunctionNew(declaration)
         }
         if (declarationParent.isObject) {
-          declarationParent.reportError("@Redacted is useless on object classes.")
+          if (!supertypeIsRedacted) {
+            declarationParent.reportError("@Redacted is useless on object classes.")
+            return super.visitFunctionNew(declaration)
+          } else if (classIsUnredacted) {
+            declarationParent.reportError("@Unredacted is useless on object classes.")
+            return super.visitFunctionNew(declaration)
+          }
+        }
+        if (classIsRedacted && classIsUnredacted) {
+          declarationParent.reportError(
+            "@Redacted and @Unredacted cannot be applied to a single class."
+          )
+          return super.visitFunctionNew(declaration)
+        }
+        if (classIsUnredacted && !supertypeIsRedacted) {
+          declarationParent.reportError(
+            "@Unredacted cannot be applied to a class unless a supertype is marked @Redacted."
+          )
           return super.visitFunctionNew(declaration)
         }
         if (anyUnredacted && (!classIsRedacted && !supertypeIsRedacted)) {
@@ -149,7 +167,12 @@ internal class RedactedIrVisitor(
           )
           return super.visitFunctionNew(declaration)
         }
-        declaration.convertToGeneratedToString(properties, classIsRedacted, supertypeIsRedacted)
+        declaration.convertToGeneratedToString(
+          properties,
+          classIsRedacted,
+          classIsUnredacted,
+          supertypeIsRedacted,
+        )
       }
     }
 
@@ -166,6 +189,7 @@ internal class RedactedIrVisitor(
   private fun IrFunction.convertToGeneratedToString(
     properties: List<Property>,
     classIsRedacted: Boolean,
+    classIsUnredacted: Boolean,
     supertypeIsRedacted: Boolean,
   ) {
     val parent = parent as IrClass
@@ -179,6 +203,7 @@ internal class RedactedIrVisitor(
           irFunction = this@convertToGeneratedToString,
           irProperties = properties,
           classIsRedacted = classIsRedacted,
+          classIsUnredacted = classIsUnredacted,
           supertypeIsRedacted = supertypeIsRedacted,
         )
       }
@@ -209,12 +234,13 @@ internal class RedactedIrVisitor(
     irFunction: IrFunction,
     irProperties: List<Property>,
     classIsRedacted: Boolean,
+    classIsUnredacted: Boolean,
     supertypeIsRedacted: Boolean,
   ) {
     val irConcat = irConcat()
     irConcat.addArgument(irString(irClass.name.asString() + "("))
     val hasUnredactedProperties by lazy(NONE) { irProperties.any { it.isUnredacted } }
-    if (classIsRedacted && !hasUnredactedProperties) {
+    if (classIsRedacted && !classIsUnredacted && !hasUnredactedProperties) {
       irConcat.addArgument(irString(replacementString))
     } else {
       var first = true
@@ -225,7 +251,7 @@ internal class RedactedIrVisitor(
         val redactProperty =
           property.isRedacted ||
             (classIsRedacted && !property.isUnredacted) ||
-            (supertypeIsRedacted && !property.isUnredacted)
+            (supertypeIsRedacted && !classIsUnredacted && !property.isUnredacted)
         if (redactProperty) {
           irConcat.addArgument(irString(replacementString))
         } else {

diff --git a/sample/src/jvmTest/kotlin/dev/zacsweers/redacted/sample/SmokeTest.kt b/sample/src/jvmTest/kotlin/dev/zacsweers/redacted/sample/SmokeTest.kt
@@ -34,12 +34,27 @@ class SmokeTest {
     assertThat(notSoSecretChild.toString()).isEqualTo("NotSoSecretChild(unredacted=public)")
   }
 
+  @Test
+  fun unredactedClassAbstractExample() {
+    val notAtAllSecretChild = AbstractBase.NotAtAllSecretChild("public")
+    assertThat(notAtAllSecretChild.toString()).isEqualTo("NotAtAllSecretChild(unredacted=public)")
+  }
+
+  @Test
+  fun redactedObjectAbstractExample() {
+    assertThat(AbstractBase.ProplessChild.toString()).isEqualTo("ProplessChild()")
+  }
+
   @Redacted
   abstract class AbstractBase {
 
     data class SecretChild(val redact: String) : AbstractBase()
 
     data class NotSoSecretChild(@Unredacted val unredacted: String) : AbstractBase()
+
+    @Unredacted data class NotAtAllSecretChild(val unredacted: String) : AbstractBase()
+
+    data object ProplessChild : AbstractBase()
   }
 
   @Test
@@ -54,12 +69,27 @@ class SmokeTest {
     assertThat(notSoSecretChild.toString()).isEqualTo("NotSoSecretChild(unredacted=public)")
   }
 
+  @Test
+  fun unredactedClassSealedExample() {
+    val notAtAllSecretChild = SecretParent.NotAtAllSecretChild("public")
+    assertThat(notAtAllSecretChild.toString()).isEqualTo("NotAtAllSecretChild(unredacted=public)")
+  }
+
+  @Test
+  fun redactedObjectSealedExample() {
+    assertThat(SecretParent.ProplessChild.toString()).isEqualTo("ProplessChild()")
+  }
+
   @Redacted
   sealed class SecretParent {
 
     data class SecretChild(val redact: String) : SecretParent()
 
     data class NotSoSecretChild(@Unredacted val unredacted: String) : SecretParent()
+
+    @Unredacted data class NotAtAllSecretChild(val unredacted: String) : AbstractBase()
+
+    data object ProplessChild : SecretParent()
   }
 
   @Test