Implement UnmarshalYAML for ReturnTypeSpec

pkg/codegen/schema/schema.go

@@ -1564,9 +1564,6 @@ func (returnTypeSpec *ReturnTypeSpec) UnmarshalJSON(inputJSON []byte) error {
 	}
 
 	if len(data) == 0 {
-		// then the schema supplied "outputs":{}
-		// treat as of outputs was not defined
-		returnTypeSpec = nil
 		return nil
 	}
 
@@ -1587,6 +1584,33 @@ func (returnTypeSpec *ReturnTypeSpec) UnmarshalJSON(inputJSON []byte) error {
 	return nil
 }
 
+func (returnTypeSpec *ReturnTypeSpec) UnmarshalYAML(node *yaml.Node) error {
+	var data map[string]interface{}
+	if err := node.Decode(&data); err != nil {
+		return err
+	}
+
+	if len(data) == 0 {
+		return nil
+	}
+
+	var objectSpec *ObjectTypeSpec
+	var typeSpec *TypeSpec
+	if _, hasProperties := data["properties"]; hasProperties {
+		if err := node.Decode(&objectSpec); err != nil {
+			return err
+		}
+	} else {
+		if err := node.Decode(&typeSpec); err != nil {
+			return err
+		}
+	}
+
+	returnTypeSpec.TypeSpec = typeSpec
+	returnTypeSpec.ObjectTypeSpec = objectSpec
+	return nil
+}
+
 // FunctionSpec is the serializable form of a function description.
 type FunctionSpec struct {
 	// Description is the description of the function, if any.