From 96f08ed3869bbd435808a99130f77e17afd9180a Mon Sep 17 00:00:00 2001
From: Ethan Reesor <ethan.reesor@gmail.com>
Date: Wed, 14 Apr 2021 15:06:23 -0500
Subject: [PATCH] Add option to produce array of results for blocks

---
 interp/interp.go           | 37 +++++++++++++++++++++++++++++++++++++
 interp/interp_eval_test.go |  9 +++++++++
 2 files changed, 46 insertions(+)

diff --git a/interp/interp.go b/interp/interp.go
index 9d669436e..c40c8d709 100644
--- a/interp/interp.go
+++ b/interp/interp.go
@@ -131,6 +131,8 @@ type opt struct {
 	stdin    io.Reader     // standard input
 	stdout   io.Writer     // standard output
 	stderr   io.Writer     // standard error
+
+	blockStmtAsArray bool
 }
 
 // Interpreter contains global resources and state.
@@ -244,6 +246,8 @@ type Options struct {
 	// They default to os.Stding, os.Stdout and os.Stderr respectively.
 	Stdin          io.Reader
 	Stdout, Stderr io.Writer
+
+	BlockStatementAsArray bool
 }
 
 // New returns a new interpreter.
@@ -294,6 +298,8 @@ func New(options Options) *Interpreter {
 
 	// fastChan disables the cancellable version of channel operations in evalWithContext
 	i.opt.fastChan, _ = strconv.ParseBool(os.Getenv("YAEGI_FAST_CHAN"))
+
+	i.opt.blockStmtAsArray = options.BlockStatementAsArray
 	return &i
 }
 
@@ -582,6 +588,12 @@ func (interp *Interpreter) eval(src, name string, inc bool) (res reflect.Value,
 	for _, n := range initNodes {
 		interp.run(n, interp.frame)
 	}
+
+	if interp.blockStmtAsArray && root.kind == blockStmt {
+		v := blockStmtAsArray(root, interp.frame)
+		return reflect.ValueOf(v), err
+	}
+
 	v := genValue(root)
 	res = v(interp.frame)
 
@@ -595,6 +607,31 @@ func (interp *Interpreter) eval(src, name string, inc bool) (res reflect.Value,
 	return res, err
 }
 
+// BlockResult is the result of a block statement.
+type BlockResult []interface{}
+
+func blockStmtAsArray(n *node, f *frame) BlockResult {
+	res := make(BlockResult, len(n.child))
+	for i, child := range n.child {
+		if child.kind == blockStmt {
+			res[i] = blockStmtAsArray(child, f)
+			continue
+		}
+
+		v := genValue(child)
+		rv := v(f)
+		if !rv.IsValid() {
+			continue
+		}
+
+		res[i] = rv.Interface()
+		if n, ok := res[i].(*node); ok {
+			res[i] = genFunctionWrapper(n)(f)
+		}
+	}
+	return res
+}
+
 // EvalWithContext evaluates Go code represented as a string. It returns
 // a map on current interpreted package exported symbols.
 func (interp *Interpreter) EvalWithContext(ctx context.Context, src string) (reflect.Value, error) {
diff --git a/interp/interp_eval_test.go b/interp/interp_eval_test.go
index f24f21f76..bfb28b7d9 100644
--- a/interp/interp_eval_test.go
+++ b/interp/interp_eval_test.go
@@ -1490,3 +1490,12 @@ func TestREPLDivision(t *testing.T) {
 		t.Fatal("timeout")
 	}
 }
+
+func TestEvalBlockAsArray(t *testing.T) {
+	i := interp.New(interp.Options{BlockStatementAsArray: true})
+	runTests(t, i, []testCase{
+		{src: "a := 5", res: "[5]"},
+		{src: "a = 1; b := 2", res: "[1 2]"},
+		{src: "a = 3; {b = 4}", res: "[3 [4]]"},
+	})
+}