zones: turn Zones.WriteEnv() into Zones.Env().WriteTo()

Signed-off-by: Alejandro Mery <amery@jpi.io>

cmd/jpictl/env.go

@@ -16,7 +16,8 @@ var envCmd = &cobra.Command{
 			return err
 		}
 
-		return m.WriteEnv(os.Stdout)
+		_, err = m.Env().WriteTo(os.Stdout)
+		return err
 	},
 }
 

pkg/zones/env.go

@@ -7,8 +7,20 @@ import (
 	"strings"
 )
 
-// WriteEnv generates environment variables for shell scripts
-func (m *Zones) WriteEnv(w io.Writer) error {
+// Env is a shell environment factory for this cluster
+type Env struct {
+	ZoneIterator
+}
+
+// Env returns a shell environment factory
+func (m *Zones) Env() *Env {
+	return &Env{
+		ZoneIterator: m,
+	}
+}
+
+// WriteTo generates environment variables for shell scripts
+func (m *Env) WriteTo(w io.Writer) (int64, error) {
 	var buf bytes.Buffer
 
 	m.writeEnvVarFn(&buf, genEnvZones, "ZONES")
@@ -17,11 +29,10 @@ func (m *Zones) WriteEnv(w io.Writer) error {
 		return false
 	})
 
-	_, err := buf.WriteTo(w)
-	return err
+	return buf.WriteTo(w)
 }
 
-func (m *Zones) writeEnvZone(w io.Writer, z *Zone) {
+func (m *Env) writeEnvZone(w io.Writer, z *Zone) {
 	zoneID := z.ID
 
 	// ZONE{zoneID}
@@ -42,7 +53,7 @@ func (m *Zones) writeEnvZone(w io.Writer, z *Zone) {
 	}
 }
 
-func (m *Zones) writeEnvVarFn(w io.Writer, fn func(*Zones) string, name string, args ...any) {
+func (m *Env) writeEnvVarFn(w io.Writer, fn func(*Env) string, name string, args ...any) {
 	var value string
 
 	if fn != nil {
@@ -52,7 +63,7 @@ func (m *Zones) writeEnvVarFn(w io.Writer, fn func(*Zones) string, name string,
 	m.writeEnvVar(w, value, name, args...)
 }
 
-func (*Zones) writeEnvVar(w io.Writer, value string, name string, args ...any) {
+func (*Env) writeEnvVar(w io.Writer, value string, name string, args ...any) {
 	if len(args) > 0 {
 		name = fmt.Sprintf(name, args...)
 	}
@@ -64,11 +75,13 @@ func (*Zones) writeEnvVar(w io.Writer, value string, name string, args ...any) {
 	}
 }
 
-func genEnvZones(m *Zones) string {
-	s := make([]string, 0, len(m.Zones))
-	for _, z := range m.Zones {
+func genEnvZones(m *Env) string {
+	var s []string
+
+	m.ForEachZone(func(z *Zone) bool {
 		s = append(s, fmt.Sprintf("%v", z.ID))
-	}
+		return false
+	})
 
 	return strings.Join(s, " ")
 }