Expose Function stream requirements

Add Function options for declaring stdin and stdout stream requirements,
so Function stages can participate in file-preference negotiation.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: 2 people

pipe/function.go

@@ -33,20 +33,32 @@ type StageFunc func(ctx context.Context, env Env, stdin io.Reader, stdout io.Wri
 // FunctionOption configures a Function stage.
 type FunctionOption func(*goStage)
 
+// WithStdinRequirement returns a FunctionOption declaring the stage's stdin
+// requirement.
+func WithStdinRequirement(requirement StreamRequirement) FunctionOption {
+	return func(s *goStage) {
+		s.requirements.Stdin = requirement
+	}
+}
+
+// WithStdoutRequirement returns a FunctionOption declaring the stage's stdout
+// requirement.
+func WithStdoutRequirement(requirement StreamRequirement) FunctionOption {
+	return func(s *goStage) {
+		s.requirements.Stdout = requirement
+	}
+}
+
 // ForbidStdin returns a FunctionOption declaring that the stage must not be
 // connected to stdin.
 func ForbidStdin() FunctionOption {
-	return func(s *goStage) {
-		s.requirements.Stdin = StreamForbidden
-	}
+	return WithStdinRequirement(StreamForbidden)
 }
 
 // ForbidStdout returns a FunctionOption declaring that the stage must not be
 // connected to stdout.
 func ForbidStdout() FunctionOption {
-	return func(s *goStage) {
-		s.requirements.Stdout = StreamForbidden
-	}
+	return WithStdoutRequirement(StreamForbidden)
 }
 
 // Function returns a pipeline `Stage` that will run a `StageFunc` in

pipe/pipeline_test.go

@@ -968,6 +968,24 @@ func TestFunctionOptionsForbidStreams(t *testing.T) {
 	})
 }
 
+func TestFunctionOptionsSetStreamRequirements(t *testing.T) {
+	t.Parallel()
+
+	stage := pipe.Function(
+		"file-preferring",
+		func(_ context.Context, _ pipe.Env, _ io.Reader, _ io.Writer) error {
+			return nil
+		},
+		pipe.WithStdinRequirement(pipe.StreamPreferFile),
+		pipe.WithStdoutRequirement(pipe.StreamPreferFile),
+	)
+
+	assert.Equal(t, pipe.StageRequirements{
+		Stdin:  pipe.StreamPreferFile,
+		Stdout: pipe.StreamPreferFile,
+	}, stage.Requirements())
+}
+
 func TestStreamForbiddenStdin(t *testing.T) {
 	t.Parallel()
 	ctx := context.Background()