Merge 683b3bb into 9744d11

abs-lang · Nov 27, 2019 · c827c41 · c827c41
2 parents 9744d11 + 683b3bb
commit c827c41
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 6 deletions.
diff --git a/evaluator/evaluator_test.go b/evaluator/evaluator_test.go
@@ -835,6 +835,8 @@ func TestBuiltinFunctions(t *testing.T) {
 		{`len([])`, 0},
 		{`echo("hello", "world!")`, nil},
 		{`env("CONTEXT")`, "abs"},
+		{`env("FOO")`, ""},
+		{`env("FOO", "bar")`, "bar"},
 		{`type("SOME")`, "STRING"},
 		{`type(1)`, "NUMBER"},
 		{`type({})`, "HASH"},

diff --git a/evaluator/functions.go b/evaluator/functions.go
@@ -116,9 +116,9 @@ func getFns() map[string]*object.Builtin {
 			Types: []string{},
 			Fn:    stdinFn,
 		},
-		// env(variable:"PWD")
+		// env(variable:"PWD") or env(string:"KEY", string:"VAL")
 		"env": &object.Builtin{
-			Types: []string{object.STRING_OBJ},
+			Types: []string{},
 			Fn:    envFn,
 		},
 		// arg(position:1)
@@ -374,6 +374,22 @@ func validateArgs(tok token.Token, name string, args []object.Object, size int,
 	return nil
 }
 
+func validateVarArgs(tok token.Token, name string, args []object.Object, required int, types [][][]string) object.Object {
+	if len(args) < required {
+		return newError(tok, "wrong number of arguments to %s(...): got=%d, min=%d, max=%d", name, len(args), required, len(types))
+	}
+
+	for i, set := range types {
+		for _, t := range set {
+			if !util.Contains(t, string(args[i].Type())) {
+				return newError(tok, "argument %d to %s(...) is not supported (got: %s, allowed: %s)", i, name, args[i].Inspect(), strings.Join(t, ", "))
+			}
+		}
+	}
+
+	return nil
+}
+
 // len(var:"hello")
 func lenFn(tok token.Token, env *object.Environment, args ...object.Object) object.Object {
 	err := validateArgs(tok, "len", args, 1, [][]string{{object.STRING_OBJ, object.ARRAY_OBJ}})
@@ -704,15 +720,21 @@ func stdinNextFn() (object.Object, object.Object) {
 	return &object.Number{Value: float64(scannerPosition)}, &object.String{Token: tok, Value: scanner.Text()}
 }
 
-// env(variable:"PWD")
+// env(variable:"PWD") or env(string:"KEY", string:"VAL")
 func envFn(tok token.Token, env *object.Environment, args ...object.Object) object.Object {
-	err := validateArgs(tok, "env", args, 1, [][]string{{object.STRING_OBJ}})
+	err := validateVarArgs(tok, "env", args, 1, [][][]string{{{object.STRING_OBJ}, {object.STRING_OBJ}}})
 	if err != nil {
 		return err
 	}
 
-	arg := args[0].(*object.String)
-	return &object.String{Token: tok, Value: os.Getenv(arg.Value)}
+	key := args[0].(*object.String)
+
+	if len(args) > 1 {
+		val := args[1].(*object.String)
+		os.Setenv(key.Value, val.Value)
+	}
+
+	return &object.String{Token: tok, Value: os.Getenv(key.Value)}
 }
 
 // arg(position:1)