From 5d8e0e0fdd5e1354e365bab98d31c21c58717ea5 Mon Sep 17 00:00:00 2001
From: David Alberto Adler <dalberto.adler@gmail.com>
Date: Mon, 1 Jul 2024 10:58:43 +0100
Subject: [PATCH] Add proper marshalling of overlays

---
 workflow/source.go | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/workflow/source.go b/workflow/source.go
index 7a3ba4c..4e14836 100644
--- a/workflow/source.go
+++ b/workflow/source.go
@@ -26,6 +26,34 @@ type Overlay struct {
 	Document            *Document            `yaml:"document,omitempty"`
 }
 
+func (o *Overlay) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	var doc Document
+	if err := unmarshal(&doc); err == nil {
+		o.Document = &doc
+		return nil
+	}
+
+	var fallback FallbackCodeSamples
+	if err := unmarshal(&fallback); err == nil {
+		o.FallbackCodeSamples = &fallback
+		return nil
+	}
+
+	return fmt.Errorf("failed to unmarshal Overlay")
+}
+
+func (o Overlay) MarshalYAML() (interface{}, error) {
+	if o.Document != nil {
+		return o.Document, nil
+	}
+
+	if o.FallbackCodeSamples != nil {
+		return o.FallbackCodeSamples, nil
+	}
+
+	return nil, fmt.Errorf("failed to marshal Overlay")
+}
+
 type FallbackCodeSamples struct {
 	FallbackCodeSamplesLanguage string `yaml:"fallbackCodeSamplesLanguage,omitempty"`
 }