Anton's implementation for removeAddedFields

Co-authored-by: Anton Parkhomenko <chuwy@users.noreply.github.com>
snowplow · Sep 23, 2020 · 2d5f6a5 · 2d5f6a5
1 parent cb3b570
commit 2d5f6a5
Showing 1 changed file with 17 additions and 29 deletions.
diff --git a/...nalytics.snowplow.enrich/common/enrichments/registry/pii/PiiPseudonymizerEnrichment.scala b/...nalytics.snowplow.enrich/common/enrichments/registry/pii/PiiPseudonymizerEnrichment.scala
@@ -135,35 +135,23 @@ object PiiPseudonymizerEnrichment extends ParseableEnrichment {
       .getOrElse(s"The specified json field $fieldName is not supported".asLeft)
 
   /** Helper to remove fields that were wrongly added and are not in the original JSON. See #351. */
-  private[pii] def removeAddedFields(hashed: Json, original: Json): Json =
-    hashed.asObject.map(_.toMap) match {
-      case Some(hashed_fields) => // hashed is JSON object
-        original.asObject.map(_.toMap) match {
-          case Some(orig_fields) =>
-            val newMap =
-              hashed_fields
-                .collect {
-                  case (k, v) if orig_fields.isDefinedAt(k) =>
-                    (k, removeAddedFields(v, orig_fields.get(k).get))
-                }
-            Json.fromFields(newMap)
-          case None =>
-            hashed // should never happen. Would mean change of type of one field
-        }
-      case None =>
-        hashed.asArray match {
-          case Some(hashed_arr) => // hashed is array (can contain JSON objects)
-            original.asArray match {
-              case Some(orig_arr) =>
-                // we can use zip because there should never be new fields in an array, only in an object
-                val values = hashed_arr.zip(orig_arr).map { case (hashed, orig) => removeAddedFields(hashed, orig) }
-                Json.fromValues(values)
-              case None =>
-                Json.fromValues(hashed_arr) // should never happen. Would mean change of type of one field
-            }
-          case None => hashed // hashed is neither JSON object nor array
-        }
-    }
+  private[pii] def removeAddedFields(hashed: Json, original: Json): Json = {
+    val fixedObject = for {
+      hashedFields <- hashed.asObject
+      originalFields <- original.asObject
+      newFields = hashedFields.toList.flatMap {
+        case (k, v) => originalFields(k).map(origV => (k, removeAddedFields(v, origV)))
+      }
+    } yield Json.fromFields(newFields)
+
+    lazy val fixedArray = for {
+      hashedArr <- hashed.asArray
+      originalArr <- original.asArray
+      newArr = hashedArr.zip(originalArr).map { case (hashed, orig) => removeAddedFields(hashed, orig) }
+    } yield Json.fromValues(newArr)
+
+    fixedObject.orElse(fixedArray).getOrElse(hashed)
+  }
 }
 
 /**