Support formatting options when unparsing ASTs to strings

common/operators/operators.go

@@ -141,3 +141,13 @@ func Precedence(symbol string) int {
 	}
 	return op.precedence
 }
+
+// Arity returns the number of argument the operator takes
+// -1 is returned if an undefined symbol is provided
+func Arity(symbol string) int {
+	op, found := operatorMap[symbol]
+	if !found {
+		return -1
+	}
+	return op.arity
+}

parser/unparser.go

@@ -36,9 +36,29 @@ import (
 // - Floating point values are converted to the small number of digits needed to represent the value.
 // - Spacing around punctuation marks may be lost.
 // - Parentheses will only be applied when they affect operator precedence.
-func Unparse(expr *exprpb.Expr, info *exprpb.SourceInfo) (string, error) {
-	un := &unparser{info: info}
-	err := un.visit(expr)
+//
+// This function optionally takes in one or more UnparserOption to alter the unparsing behavior, such as
+// performing word wrapping on expressions.
+func Unparse(expr *exprpb.Expr, info *exprpb.SourceInfo, opts ...UnparserOption) (string, error) {
+	unparserOpts := &unparserOption{
+		wrapOnColumn:         defaultWrapOnColumn,
+		wrapAfterColumnLimit: defaultWrapAfterColumnLimit,
+		operatorsToWrapOn:    defaultOperatorsToWrapOn,
+	}
+
+	var err error
+	for _, opt := range opts {
+		unparserOpts, err = opt(unparserOpts)
+		if err != nil {
+			return "", err
+		}
+	}
+
+	un := &unparser{
+		info:    info,
+		options: unparserOpts,
+	}
+	err = un.visit(expr)
 	if err != nil {
 		return "", err
 	}
@@ -47,8 +67,10 @@ func Unparse(expr *exprpb.Expr, info *exprpb.SourceInfo) (string, error) {
 
 // unparser visits an expression to reconstruct a human-readable string from an AST.
 type unparser struct {
-	str  strings.Builder
-	info *exprpb.SourceInfo
+	str              strings.Builder
+	info             *exprpb.SourceInfo
+	options          *unparserOption
+	lastWrappedIndex int
 }
 
 func (un *unparser) visit(expr *exprpb.Expr) error {
@@ -135,9 +157,8 @@ func (un *unparser) visitCallBinary(expr *exprpb.Expr) error {
 	if !found {
 		return fmt.Errorf("cannot unmangle operator: %s", fun)
 	}
-	un.str.WriteString(" ")
-	un.str.WriteString(unmangled)
-	un.str.WriteString(" ")
+
+	un.writeOperatorWithWrapping(fun, unmangled)
 	return un.visitMaybeNested(rhs, rhsParen)
 }
 
@@ -151,14 +172,16 @@ func (un *unparser) visitCallConditional(expr *exprpb.Expr) error {
 	if err != nil {
 		return err
 	}
-	un.str.WriteString(" ? ")
+	un.writeOperatorWithWrapping(operators.Conditional, "?")
+
 	// add parens if operand is a conditional itself.
 	nested = isSamePrecedence(operators.Conditional, args[1]) ||
 		isComplexOperator(args[1])
 	err = un.visitMaybeNested(args[1], nested)
 	if err != nil {
 		return err
 	}
+
 	un.str.WriteString(" : ")
 	// add parens if operand is a conditional itself.
 	nested = isSamePrecedence(operators.Conditional, args[2]) ||
@@ -444,3 +467,130 @@ func bytesToOctets(byteVal []byte) string {
 	}
 	return b.String()
 }
+
+// writeOperatorWithWrapping outputs the operator and inserts a newline for operators configured
+// in the unparser options.
+func (un *unparser) writeOperatorWithWrapping(fun string, unmangled string) bool {
+	_, wrapOperatorExists := un.options.operatorsToWrapOn[fun]
+	lineLength := un.str.Len() - un.lastWrappedIndex + len(fun)
+
+	if wrapOperatorExists && lineLength >= un.options.wrapOnColumn {
+		un.lastWrappedIndex = un.str.Len()
+		// wrapAfterColumnLimit flag dictates whether the newline is placed
+		// before or after the operator
+		if un.options.wrapAfterColumnLimit {
+			// Input: a && b
+			// Output: a &&\nb
+			un.str.WriteString(" ")
+			un.str.WriteString(unmangled)
+			un.str.WriteString("\n")
+		} else {
+			// Input: a && b
+			// Output: a\n&& b
+			un.str.WriteString("\n")
+			un.str.WriteString(unmangled)
+			un.str.WriteString(" ")
+		}
+		return true
+	} else {
+		un.str.WriteString(" ")
+		un.str.WriteString(unmangled)
+		un.str.WriteString(" ")
+	}
+	return false
+}
+
+// Defined defaults for the unparser options
+var (
+	defaultWrapOnColumn         = 80
+	defaultWrapAfterColumnLimit = true
+	defaultOperatorsToWrapOn    = map[string]bool{
+		operators.LogicalAnd: true,
+		operators.LogicalOr:  true,
+	}
+)
+
+// UnparserOption is a functional option for configuring the output formatting
+// of the Unparse function.
+type UnparserOption func(*unparserOption) (*unparserOption, error)
+
+// Internal representation of the UnparserOption type
+type unparserOption struct {
+	wrapOnColumn         int
+	operatorsToWrapOn    map[string]bool
+	wrapAfterColumnLimit bool
+}
+
+// WrapOnColumn wraps the output expression when its string length exceeds a specified limit
+// for operators set by WrapOnOperators function or by default, "&&" and "||" will be wrapped.
+//
+// Example usage:
+//
+//	Unparse(expr, sourceInfo, WrapOnColumn(40), WrapOnOperators(Operators.LogicalAnd))
+//
+// This will insert a newline immediately after the logical AND operator for the below example input:
+//
+// Input:
+// 'my-principal-group' in request.auth.claims && request.auth.claims.iat > now - duration('5m')
+//
+// Output:
+// 'my-principal-group' in request.auth.claims &&
+// request.auth.claims.iat > now - duration('5m')
+func WrapOnColumn(col int) UnparserOption {
+	return func(opt *unparserOption) (*unparserOption, error) {
+		if col < 1 {
+			return nil, fmt.Errorf("Invalid unparser option. Wrap column value must be greater than or equal to 1. Got %v instead", col)
+		}
+		opt.wrapOnColumn = col
+		return opt, nil
+	}
+}
+
+// WrapOnOperators specifies which operators to perform word wrapping on an output expression when its string length
+// exceeds the column limit set by WrapOnColumn function.
+//
+// Word wrapping is supported on non-unary symbolic operators. Refer to operators.go for the full list
+//
+// This will replace any previously supplied operators instead of merging them.
+func WrapOnOperators(symbols ...string) UnparserOption {
+	return func(opt *unparserOption) (*unparserOption, error) {
+		opt.operatorsToWrapOn = make(map[string]bool)
+		for _, symbol := range symbols {
+			_, found := operators.FindReverse(symbol)
+			if !found {
+				return nil, fmt.Errorf("Invalid unparser option. Unsupported operator: %s", symbol)
+			}
+			arity := operators.Arity(symbol)
+			if arity < 2 {
+				return nil, fmt.Errorf("Invalid unparser option. Unary operators are unsupported: %s", symbol)
+			}
+
+			opt.operatorsToWrapOn[symbol] = true
+		}
+
+		return opt, nil
+	}
+}
+
+// WrapAfterColumnLimit dictates whether to insert a newline before or after the specified operator
+// when word wrapping is performed.
+//
+// Example usage:
+//
+//	Unparse(expr, sourceInfo, WrapOnColumn(40), WrapOnOperators(Operators.LogicalAnd), WrapAfterColumnLimit(false))
+//
+// This will insert a newline immediately before the logical AND operator for the below example input, ensuring
+// that the length of a line never exceeds the specified column limit:
+//
+// Input:
+// 'my-principal-group' in request.auth.claims && request.auth.claims.iat > now - duration('5m')
+//
+// Output:
+// 'my-principal-group' in request.auth.claims
+// && request.auth.claims.iat > now - duration('5m')
+func WrapAfterColumnLimit(wrapAfter bool) UnparserOption {
+	return func(opt *unparserOption) (*unparserOption, error) {
+		opt.wrapAfterColumnLimit = wrapAfter
+		return opt, nil
+	}
+}