Add capability to unparse optionals

PiperOrigin-RevId: 566700621
google · Sep 19, 2023 · af584c2 · af584c2
1 parent d39039e
commit af584c2
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 14 deletions.
diff --git a/parser/src/main/java/dev/cel/parser/CelUnparserVisitor.java b/parser/src/main/java/dev/cel/parser/CelUnparserVisitor.java
@@ -27,6 +27,7 @@
 import dev.cel.common.ast.CelExpr.CelSelect;
 import dev.cel.common.ast.CelExpr.ExprKind.Kind;
 import dev.cel.common.ast.CelExprVisitor;
+import java.util.HashSet;
 import java.util.Optional;
 
 /** Visitor implementation to unparse an AST. */
@@ -104,16 +105,7 @@ protected void visit(CelExpr expr, CelIdent ident) {
 
   @Override
   protected void visit(CelExpr expr, CelSelect select) {
-    if (select.testOnly()) {
-      stringBuilder.append(Operator.HAS.getFunction()).append(LEFT_PAREN);
-    }
-    CelExpr operand = select.operand();
-    boolean nested = !select.testOnly() && isBinaryOrTernaryOperator(operand);
-    visitMaybeNested(operand, nested);
-    stringBuilder.append(DOT).append(select.field());
-    if (select.testOnly()) {
-      stringBuilder.append(RIGHT_PAREN);
-    }
+    visitSelect(select.operand(), select.testOnly(), DOT, select.field());
   }
 
   @Override
@@ -133,7 +125,12 @@ protected void visit(CelExpr expr, CelCall call) {
     }
 
     if (fun.equals(Operator.INDEX.getFunction())) {
-      visitIndex(call);
+      visitIndex(call, LEFT_BRACKET);
+      return;
+    }
+
+    if (fun.equals(Operator.OPTIONAL_INDEX.getFunction())) {
+      visitIndex(call, LEFT_BRACKET + QUESTION_MARK);
       return;
     }
 
@@ -142,6 +139,13 @@ protected void visit(CelExpr expr, CelCall call) {
       return;
     }
 
+    if (fun.equals(Operator.OPTIONAL_SELECT.getFunction())) {
+      CelExpr operand = call.args().get(0);
+      String field = call.args().get(1).constant().stringValue();
+      visitSelect(operand, false, DOT + QUESTION_MARK, field);
+      return;
+    }
+
     if (call.target().isPresent()) {
       boolean nested = isBinaryOrTernaryOperator(call.target().get());
       visitMaybeNested(call.target().get(), nested);
@@ -161,10 +165,14 @@ protected void visit(CelExpr expr, CelCall call) {
   @Override
   protected void visit(CelExpr expr, CelCreateList createList) {
     stringBuilder.append(LEFT_BRACKET);
+    HashSet<Integer> optionalIndices = new HashSet<>(createList.optionalIndices());
     for (int i = 0; i < createList.elements().size(); i++) {
       if (i > 0) {
         stringBuilder.append(COMMA).append(SPACE);
       }
+      if (optionalIndices.contains(i)) {
+        stringBuilder.append(QUESTION_MARK);
+      }
       visit(createList.elements().get(i));
     }
     stringBuilder.append(RIGHT_BRACKET);
@@ -180,6 +188,9 @@ protected void visit(CelExpr expr, CelCreateStruct createStruct) {
       }
 
       CelCreateStruct.Entry e = createStruct.entries().get(i);
+      if (e.optionalEntry()) {
+        stringBuilder.append(QUESTION_MARK);
+      }
       stringBuilder.append(e.fieldKey());
       stringBuilder.append(COLON).append(SPACE);
       visit(e.value());
@@ -196,6 +207,9 @@ protected void visit(CelExpr expr, CelCreateMap createMap) {
       }
 
       CelCreateMap.Entry e = createMap.entries().get(i);
+      if (e.optionalEntry()) {
+        stringBuilder.append(QUESTION_MARK);
+      }
       visit(e.key());
       stringBuilder.append(COLON).append(SPACE);
       visit(e.value());
@@ -243,6 +257,18 @@ private void visitBinary(CelCall expr, String op) {
     visitMaybeNested(rhs, rhsParen);
   }
 
+  private void visitSelect(CelExpr operand, boolean testOnly, String op, String field) {
+    if (testOnly) {
+      stringBuilder.append(Operator.HAS.getFunction()).append(LEFT_PAREN);
+    }
+    boolean nested = !testOnly && isBinaryOrTernaryOperator(operand);
+    visitMaybeNested(operand, nested);
+    stringBuilder.append(op).append(field);
+    if (testOnly) {
+      stringBuilder.append(RIGHT_PAREN);
+    }
+  }
+
   private void visitTernary(CelCall expr) {
     if (expr.args().size() != 3) {
       throw new IllegalArgumentException(String.format("unexpected ternary: %s", expr));
@@ -268,13 +294,13 @@ private void visitTernary(CelCall expr) {
     visitMaybeNested(expr.args().get(2), nested);
   }
 
-  private void visitIndex(CelCall expr) {
+  private void visitIndex(CelCall expr, String op) {
     if (expr.args().size() != 2) {
       throw new IllegalArgumentException(String.format("unexpected index call: %s", expr));
     }
     boolean nested = isBinaryOrTernaryOperator(expr.args().get(0));
     visitMaybeNested(expr.args().get(0), nested);
-    stringBuilder.append(LEFT_BRACKET);
+    stringBuilder.append(op);
     visit(expr.args().get(1));
     stringBuilder.append(RIGHT_BRACKET);
   }

diff --git a/parser/src/test/java/dev/cel/parser/CelUnparserImplTest.java b/parser/src/test/java/dev/cel/parser/CelUnparserImplTest.java
@@ -28,6 +28,7 @@
 import dev.cel.common.ast.CelExpr.CelCall;
 import dev.cel.common.ast.CelExpr.CelCreateMap;
 import dev.cel.common.ast.CelExpr.CelCreateStruct;
+import dev.cel.extensions.CelOptionalLibrary;
 import java.util.Arrays;
 import java.util.List;
 import org.junit.Test;
@@ -39,6 +40,7 @@ public final class CelUnparserImplTest {
   private final CelParserImpl parser =
       CelParserImpl.newBuilder()
           .setOptions(CelOptions.newBuilder().populateMacroCalls(true).build())
+          .addLibraries(CelOptionalLibrary.INSTANCE)
           .addMacros(CelMacro.STANDARD_MACROS)
           .build();
 
@@ -153,7 +155,14 @@ public List<String> provideValues() {
           "size(x) == x.size()",
 
           // Long string
-          "Loooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooong");
+          "Loooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooong",
+
+          // Optionals
+          "a.?b",
+          "a[?b]",
+          "[?a, ?b, c]",
+          "{?a: b, c: d}",
+          "v1alpha1.Expr{?id: id, call_expr: v1alpha1.Call_Expr{function: \"name\"}}");
     }
   }