From 4cea33273772eee8f3aa24283dd708939047b72d Mon Sep 17 00:00:00 2001
From: mitchell <mitchellb@activestate.com>
Date: Fri, 6 Sep 2024 12:43:15 -0400
Subject: [PATCH] Added marshaling and pre-unmarshaling hooks for buildscript
 <-> buildexpression.

This allows for outside packages to hook into:
- Unmarshaling from buildexpression to buildscript
- Marshaling from buildscript to buildexpression
---
 pkg/buildscript/marshal_buildexpression.go   | 19 ++++++++++++--
 pkg/buildscript/unmarshal_buildexpression.go | 27 +++++++++++++++++++-
 2 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/pkg/buildscript/marshal_buildexpression.go b/pkg/buildscript/marshal_buildexpression.go
index d9dbe61c99..1f55f87133 100644
--- a/pkg/buildscript/marshal_buildexpression.go
+++ b/pkg/buildscript/marshal_buildexpression.go
@@ -21,6 +21,21 @@ const (
 	requirementComparatorKey          = "comparator"
 )
 
+type MarshalerFunc func([]*Value) ([]byte, error)
+
+var marshalers map[string]MarshalerFunc
+
+func init() {
+	marshalers = make(map[string]MarshalerFunc)
+	RegisterFunctionMarshaler("Req", marshalReq) // marshal into legacy object format for now
+}
+
+// RegisterFunctionMarshaler registers a buildexpression marshaler for a buildscript function.
+// Marshalers accept a buildscript Value, and marshals it to buildexpression JSON (e.g. an object).
+func RegisterFunctionMarshaler(name string, marshalJSON MarshalerFunc) {
+	marshalers[name] = marshalJSON
+}
+
 // MarshalJSON returns this structure as a build expression in JSON format, suitable for sending to
 // the Platform.
 func (b *BuildScript) MarshalBuildExpression() ([]byte, error) {
@@ -91,8 +106,8 @@ func (v *Value) MarshalJSON() ([]byte, error) {
 }
 
 func (f *FuncCall) MarshalJSON() ([]byte, error) {
-	if f.Name == reqFuncName {
-		return marshalReq(f.Arguments) // marshal into legacy object format for now
+	if marshalJSON, exists := marshalers[f.Name]; exists {
+		return marshalJSON(f.Arguments)
 	}
 
 	m := make(map[string]interface{})
diff --git a/pkg/buildscript/unmarshal_buildexpression.go b/pkg/buildscript/unmarshal_buildexpression.go
index 08e0c3b6e6..59479995b2 100644
--- a/pkg/buildscript/unmarshal_buildexpression.go
+++ b/pkg/buildscript/unmarshal_buildexpression.go
@@ -43,6 +43,22 @@ const (
 	inKey  = "in"
 )
 
+type PreUnmarshalerFunc func(map[string]interface{}) (map[string]interface{}, error)
+
+var preUnmarshalers map[string]PreUnmarshalerFunc
+
+func init() {
+	preUnmarshalers = make(map[string]PreUnmarshalerFunc)
+}
+
+// RegisterFunctionPreUnmarshaler registers a buildscript pre-unmarshaler for a buildexpression
+// function.
+// Pre-unmarshalers accept a JSON object of function arguments, transform those arguments as
+// necessary, and return a JSON object for final unmarshaling to buildscript.
+func RegisterFunctionPreUnmarshaler(name string, preUnmarshal PreUnmarshalerFunc) {
+	preUnmarshalers[name] = preUnmarshal
+}
+
 // UnmarshalBuildExpression returns a BuildScript constructed from the given build expression in
 // JSON format.
 // Build scripts and build expressions are almost identical, with the exception of the atTime field.
@@ -251,12 +267,21 @@ func unmarshalFuncCall(path []string, m map[string]interface{}) (*FuncCall, erro
 	var name string
 	var argsInterface interface{}
 	for key, value := range m {
-		_, ok := value.(map[string]interface{})
+		m, ok := value.(map[string]interface{})
 		if !ok {
 			return nil, errs.New("Incorrect argument format")
 		}
 
 		name = key
+
+		if preUnmarshal, exists := preUnmarshalers[name]; exists {
+			var err error
+			value, err = preUnmarshal(m)
+			if err != nil {
+				return nil, errs.Wrap(err, "Unable to pre-unmarshal function '%s'", name)
+			}
+		}
+
 		argsInterface = value
 	}