From 0c433401420a3f2bc56e95f072324851f224c383 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stanis=C5=82aw=20Barzowski?= <sbarzowski@google.com>
Date: Fri, 6 Oct 2017 19:16:36 -0400
Subject: [PATCH] Support for import callbacks

It was already there, it was a matter of exposing it in the API
---
 imports.go   |  6 +++---
 main_test.go | 20 ++++++++++++++++++++
 vm.go        |  8 +++++++-
 3 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/imports.go b/imports.go
index a2c7151..5bd2101 100644
--- a/imports.go
+++ b/imports.go
@@ -147,9 +147,9 @@ type MemoryImporter struct {
 	data map[string]string
 }
 
-func (importer *MemoryImporter) Import(dir, importedPath string) (*ImportedData, error) {
+func (importer *MemoryImporter) Import(dir, importedPath string) *ImportedData {
 	if content, ok := importer.data[importedPath]; ok {
-		return &ImportedData{content: content, foundHere: importedPath}, nil
+		return &ImportedData{content: content, foundHere: importedPath}
 	}
-	return nil, fmt.Errorf("Import not available %v", importedPath)
+	return &ImportedData{err: fmt.Errorf("Import not available %v", importedPath)}
 }
diff --git a/main_test.go b/main_test.go
index 63d8cdc..bf45536 100644
--- a/main_test.go
+++ b/main_test.go
@@ -195,3 +195,23 @@ func TestMinimalError(t *testing.T) {
 
 // TODO(sbarzowski) test pretty errors once they are stable-ish
 // probably "golden" pattern is the right one for that
+
+func TestCustomImporter(t *testing.T) {
+	vm := MakeVM()
+	vm.Importer(&MemoryImporter{
+		map[string]string{
+			"a.jsonnet": "2 + 2",
+			"b.jsonnet": "3 + 3",
+		},
+	})
+	input := `[import "a.jsonnet", importstr "b.jsonnet"]`
+	expected := `[ 4, "3 + 3" ]`
+	actual, err := vm.EvaluateSnippet("custom_import.jsonnet", input)
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+	actual = removeExcessiveWhitespace(actual)
+	if actual != expected {
+		t.Errorf("Expected %v, but got %v", expected, actual)
+	}
+}
diff --git a/vm.go b/vm.go
index c37cbb7..6cfbc7c 100644
--- a/vm.go
+++ b/vm.go
@@ -52,6 +52,7 @@ func MakeVM() *VM {
 		MaxStack: 500,
 		ext:      make(vmExtMap),
 		ef:       ErrorFormatter{pretty: true, colorful: true, MaxStackTraceSize: 20},
+		importer: &FileImporter{},
 	}
 }
 
@@ -65,6 +66,11 @@ func (vm *VM) ExtCode(key string, val string) {
 	vm.ext[key] = vmExt{value: val, isCode: true}
 }
 
+// Importer sets Importer to use during evaluation (import callback)
+func (vm *VM) Importer(i Importer) {
+	vm.importer = i
+}
+
 func (vm *VM) evaluateSnippet(filename string, snippet string) (output string, err error) {
 	defer func() {
 		if r := recover(); r != nil {
@@ -75,7 +81,7 @@ func (vm *VM) evaluateSnippet(filename string, snippet string) (output string, e
 	if err != nil {
 		return "", err
 	}
-	output, err = evaluate(node, vm.ext, vm.MaxStack, &FileImporter{})
+	output, err = evaluate(node, vm.ext, vm.MaxStack, vm.importer)
 	if err != nil {
 		return "", err
 	}