Add quantity library to CEL

staging/src/k8s.io/apiserver/pkg/cel/environment/base.go

@@ -78,6 +78,7 @@ var baseOpts = []VersionedOptions{
 		EnvOptions: []cel.EnvOption{
 			cel.CrossTypeNumericComparisons(true),
 			cel.OptionalTypes(),
+			library.Quantity(),
 		},
 	},
 	// TODO: switch to ext.Strings version 2 once format() is fixed to work with HomogeneousAggregateLiterals.

staging/src/k8s.io/apiserver/pkg/cel/library/quantity.go

@@ -0,0 +1,375 @@
+/*
+Copyright 2023 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package library
+
+import (
+	"errors"
+
+	"github.com/google/cel-go/cel"
+	"github.com/google/cel-go/common/types"
+	"github.com/google/cel-go/common/types/ref"
+	"k8s.io/apimachinery/pkg/api/resource"
+	apiservercel "k8s.io/apiserver/pkg/cel"
+)
+
+// Quantity provides a CEL function library extension of Kubernetes
+// resource.Quantity parsing functions. See `resource.Quantity`
+// documentation for more detailed information about the format itself:
+// https://pkg.go.dev/k8s.io/apimachinery/pkg/api/resource#Quantity
+//
+// quantity
+//
+// Converts a string to a Quantity or results in an error if the string is not a valid Quantity. Refer
+// to resource.Quantity documentation for information on accepted patterns.
+//
+//	quantity(<string>) <Quantity>
+//
+// Examples:
+//
+//	quantity('1.5G') // returns a Quantity
+//	quantity('200k') // returns a Quantity
+//	quantity('200K') // error
+//	quantity('Three') // error
+//	quantity('Mi') // error
+//
+// isQuantity
+//
+// Returns true if a string is a valid Quantity. isQuantity returns true if and
+// only if quantity does not result in error.
+//
+//	isQuantity( <string>) <bool>
+//
+// Examples:
+//
+//	isQuantity('1.3G') // returns true
+//	isQuantity('1.3Gi') // returns true
+//	isQuantity('1,3G') // returns false
+//	isQuantity('10000k') // returns true
+//	isQuantity('200K') // returns false
+//	isQuantity('Three') // returns false
+//	isQuantity('Mi') // returns false
+//
+// Conversion to Scalars:
+//
+//   - isInteger: returns true if and only if asInteger is safe to call without an error
+//
+//   - asInteger: returns a representation of the current value as an int64 if
+//     possible or results in an error if conversion would result in overflow
+//	   or loss of precision.
+//
+//   - asApproximateFloat: returns a float64 representation of the quantity which may
+//     lose precision. If the value of the quantity is outside the range of a float64
+//     +Inf/-Inf will be returned.
+//
+//     <Quantity>.isInteger() <bool>
+//     <Quantity>.asInteger() <int>
+//     <Quantity>.asApproximateFloat() <float>
+//
+// Examples:
+//
+// quantity("50000000G").isInteger() // returns true
+// quantity("50k").isInteger() // returns true
+// quantity("9999999999999999999999999999999999999G").asInteger() // error: cannot convert value to integer
+// quantity("9999999999999999999999999999999999999G").isInteger() // returns false
+// quantity("50k").asInteger() == 50000 // returns true
+// quantity("50k").sub(20000).asApproximateFloat() == 30000 // returns true
+//
+// Arithmetic
+//
+//   - sign: Returns `1` if the quantity is positive, `-1` if it is negative. `0` if it is zero
+//
+//   - add: Returns sum of two quantities or a quantity and an integer
+//
+//   - sub: Returns difference between two quantities or a quantity and an integer
+//
+//     <Quantity>.sign() <int>
+//     <Quantity>.add(<quantity>) <quantity>
+//     <Quantity>.add(<integer>) <quantity>
+//     <Quantity>.sub(<quantity>) <quantity>
+//     <Quantity>.sub(<integer>) <quantity>
+//
+// Examples:
+//
+// quantity("50k").add("20k") == quantity("70k") // returns true
+// quantity("50k").add(20) == quantity("50020") // returns true
+// quantity("50k").sub("20k") == quantity("30k") // returns true
+// quantity("50k").sub(20000) == quantity("30k") // returns true
+// quantity("50k").add(20).sub(quantity("100k")).sub(-50000) == quantity("20") // returns true
+//
+// Comparisons
+//
+//   - isGreaterThan: Returns true if and only if the receiver is greater than the operand
+//
+//   - isLessThan: Returns true if and only if the receiver is less than the operand
+//
+//   - compareTo: Compares receiver to operand and returns 0 if they are equal, 1 if the receiver is greater, or -1 if the receiver is less than the operand
+//
+//
+//     <Quantity>.isLessThan(<quantity>) <bool>
+//     <Quantity>.isGreaterThan(<quantity>) <bool>
+//     <Quantity>.compareTo(<quantity>) <int>
+//
+// Examples:
+//
+// quantity("200M").compareTo(quantity("0.2G")) // returns 0
+// quantity("50M").compareTo(quantity("50Mi")) // returns -1
+// quantity("50Mi").compareTo(quantity("50M")) // returns 1
+// quantity("150Mi").isGreaterThan(quantity("100Mi")) // returns true
+// quantity("50Mi").isGreaterThan(quantity("100Mi")) // returns false
+// quantity("50M").isLessThan(quantity("100M")) // returns true
+// quantity("100M").isLessThan(quantity("50M")) // returns false
+
+func Quantity() cel.EnvOption {
+	return cel.Lib(quantityLib)
+}
+
+var quantityLib = &quantity{}
+
+type quantity struct{}
+
+var quantityLibraryDecls = map[string][]cel.FunctionOpt{
+	"quantity": {
+		cel.Overload("string_to_quantity", []*cel.Type{cel.StringType}, apiservercel.QuantityType, cel.UnaryBinding((stringToQuantity))),
+	},
+	"isQuantity": {
+		cel.Overload("is_quantity_string", []*cel.Type{cel.StringType}, cel.BoolType, cel.UnaryBinding(isQuantity)),
+	},
+	"sign": {
+		cel.Overload("quantity_sign", []*cel.Type{apiservercel.QuantityType}, cel.IntType, cel.UnaryBinding(quantityGetSign)),
+	},
+	"isGreaterThan": {
+		cel.MemberOverload("quantity_is_greater_than", []*cel.Type{apiservercel.QuantityType, apiservercel.QuantityType}, cel.BoolType, cel.BinaryBinding(quantityIsGreaterThan)),
+	},
+	"isLessThan": {
+		cel.MemberOverload("quantity_is_less_than", []*cel.Type{apiservercel.QuantityType, apiservercel.QuantityType}, cel.BoolType, cel.BinaryBinding(quantityIsLessThan)),
+	},
+	"compareTo": {
+		cel.MemberOverload("quantity_compare_to", []*cel.Type{apiservercel.QuantityType, apiservercel.QuantityType}, cel.IntType, cel.BinaryBinding(quantityCompareTo)),
+	},
+	"asApproximateFloat": {
+		cel.MemberOverload("quantity_get_float", []*cel.Type{apiservercel.QuantityType}, cel.DoubleType, cel.UnaryBinding(quantityGetApproximateFloat)),
+	},
+	"asInteger": {
+		cel.MemberOverload("quantity_get_int", []*cel.Type{apiservercel.QuantityType}, cel.IntType, cel.UnaryBinding(quantityGetValue)),
+	},
+	"isInteger": {
+		cel.MemberOverload("quantity_is_integer", []*cel.Type{apiservercel.QuantityType}, cel.BoolType, cel.UnaryBinding(quantityCanValue)),
+	},
+	"add": {
+		cel.MemberOverload("quantity_add", []*cel.Type{apiservercel.QuantityType, apiservercel.QuantityType}, apiservercel.QuantityType, cel.BinaryBinding(quantityAdd)),
+		cel.MemberOverload("quantity_add_int", []*cel.Type{apiservercel.QuantityType, cel.IntType}, apiservercel.QuantityType, cel.BinaryBinding(quantityAddInt)),
+	},
+	"sub": {
+		cel.MemberOverload("quantity_sub", []*cel.Type{apiservercel.QuantityType, apiservercel.QuantityType}, apiservercel.QuantityType, cel.BinaryBinding(quantitySub)),
+		cel.MemberOverload("quantity_sub_int", []*cel.Type{apiservercel.QuantityType, cel.IntType}, apiservercel.QuantityType, cel.BinaryBinding(quantitySubInt)),
+	},
+}
+
+func (*quantity) CompileOptions() []cel.EnvOption {
+	options := make([]cel.EnvOption, 0, len(quantityLibraryDecls))
+	for name, overloads := range quantityLibraryDecls {
+		options = append(options, cel.Function(name, overloads...))
+	}
+	return options
+}
+
+func (*quantity) ProgramOptions() []cel.ProgramOption {
+	return []cel.ProgramOption{}
+}
+
+func isQuantity(arg ref.Val) ref.Val {
+	str, ok := arg.Value().(string)
+	if !ok {
+		return types.MaybeNoSuchOverloadErr(arg)
+	}
+
+	_, err := resource.ParseQuantity(str)
+	if err != nil {
+		return types.Bool(false)
+	}
+
+	return types.Bool(true)
+}
+
+func stringToQuantity(arg ref.Val) ref.Val {
+	str, ok := arg.Value().(string)
+	if !ok {
+		return types.MaybeNoSuchOverloadErr(arg)
+	}
+
+	q, err := resource.ParseQuantity(str)
+	if err != nil {
+		return types.WrapErr(err)
+	}
+
+	return apiservercel.Quantity{Quantity: &q}
+}
+
+func quantityGetApproximateFloat(arg ref.Val) ref.Val {
+	q, ok := arg.Value().(*resource.Quantity)
+	if !ok {
+		return types.MaybeNoSuchOverloadErr(arg)
+	}
+	return types.Double(q.AsApproximateFloat64())
+}
+
+func quantityCanValue(arg ref.Val) ref.Val {
+	q, ok := arg.Value().(*resource.Quantity)
+	if !ok {
+		return types.MaybeNoSuchOverloadErr(arg)
+	}
+	_, success := q.AsInt64()
+	return types.Bool(success)
+}
+
+func quantityGetValue(arg ref.Val) ref.Val {
+	q, ok := arg.Value().(*resource.Quantity)
+	if !ok {
+		return types.MaybeNoSuchOverloadErr(arg)
+	}
+	v, success := q.AsInt64()
+	if !success {
+		return types.WrapErr(errors.New("cannot convert value to integer"))
+	}
+	return types.Int(v)
+}
+
+func quantityGetSign(arg ref.Val) ref.Val {
+	q, ok := arg.Value().(*resource.Quantity)
+	if !ok {
+		return types.MaybeNoSuchOverloadErr(arg)
+	}
+	return types.Int(q.Sign())
+}
+
+func quantityIsGreaterThan(arg ref.Val, other ref.Val) ref.Val {
+	q, ok := arg.Value().(*resource.Quantity)
+	if !ok {
+		return types.MaybeNoSuchOverloadErr(arg)
+	}
+
+	q2, ok := other.Value().(*resource.Quantity)
+	if !ok {
+		return types.MaybeNoSuchOverloadErr(arg)
+	}
+
+	return types.Bool(q.Cmp(*q2) == 1)
+}
+
+func quantityIsLessThan(arg ref.Val, other ref.Val) ref.Val {
+	q, ok := arg.Value().(*resource.Quantity)
+	if !ok {
+		return types.MaybeNoSuchOverloadErr(arg)
+	}
+
+	q2, ok := other.Value().(*resource.Quantity)
+	if !ok {
+		return types.MaybeNoSuchOverloadErr(arg)
+	}
+
+	return types.Bool(q.Cmp(*q2) == -1)
+}
+
+func quantityCompareTo(arg ref.Val, other ref.Val) ref.Val {
+	q, ok := arg.Value().(*resource.Quantity)
+	if !ok {
+		return types.MaybeNoSuchOverloadErr(arg)
+	}
+
+	q2, ok := other.Value().(*resource.Quantity)
+	if !ok {
+		return types.MaybeNoSuchOverloadErr(arg)
+	}
+
+	return types.Int(q.Cmp(*q2))
+}
+
+func quantityAdd(arg ref.Val, other ref.Val) ref.Val {
+	q, ok := arg.Value().(*resource.Quantity)
+	if !ok {
+		return types.MaybeNoSuchOverloadErr(arg)
+	}
+
+	q2, ok := other.Value().(*resource.Quantity)
+	if !ok {
+		return types.MaybeNoSuchOverloadErr(arg)
+	}
+
+	copy := *q
+	copy.Add(*q2)
+	return &apiservercel.Quantity{
+		Quantity: &copy,
+	}
+}
+
+func quantityAddInt(arg ref.Val, other ref.Val) ref.Val {
+	q, ok := arg.Value().(*resource.Quantity)
+	if !ok {
+		return types.MaybeNoSuchOverloadErr(arg)
+	}
+
+	q2, ok := other.Value().(int64)
+	if !ok {
+		return types.MaybeNoSuchOverloadErr(arg)
+	}
+
+	q2Converted := *resource.NewQuantity(q2, resource.DecimalExponent)
+
+	copy := *q
+	copy.Add(q2Converted)
+	return &apiservercel.Quantity{
+		Quantity: &copy,
+	}
+}
+
+func quantitySub(arg ref.Val, other ref.Val) ref.Val {
+	q, ok := arg.Value().(*resource.Quantity)
+	if !ok {
+		return types.MaybeNoSuchOverloadErr(arg)
+	}
+
+	q2, ok := other.Value().(*resource.Quantity)
+	if !ok {
+		return types.MaybeNoSuchOverloadErr(arg)
+	}
+
+	copy := *q
+	copy.Sub(*q2)
+	return &apiservercel.Quantity{
+		Quantity: &copy,
+	}
+}
+
+func quantitySubInt(arg ref.Val, other ref.Val) ref.Val {
+	q, ok := arg.Value().(*resource.Quantity)
+	if !ok {
+		return types.MaybeNoSuchOverloadErr(arg)
+	}
+
+	q2, ok := other.Value().(int64)
+	if !ok {
+		return types.MaybeNoSuchOverloadErr(arg)
+	}
+
+	q2Converted := *resource.NewQuantity(q2, resource.DecimalExponent)
+
+	copy := *q
+	copy.Sub(q2Converted)
+	return &apiservercel.Quantity{
+		Quantity: &copy,
+	}
+}