Support codegen for multi-stage files

codegen/codegen.go

@@ -7,6 +7,7 @@ import (
 	"strings"
 	"text/template"
 
+	"github.com/dop251/goja"
 	"github.com/rs/zerolog/log"
 
 	"github.com/HexmosTech/gabs/v2"
@@ -25,6 +26,12 @@ type SnippetArgs struct {
 	SnippetCore string
 }
 
+var globalVm *goja.Runtime
+
+func initialize() {
+	globalVm = cmdexec.GetJSVm()
+}
+
 func PrepareHTTPSnippetGenerator(snippetArgs SnippetArgs) string {
 	var templOutput bytes.Buffer
 	templStr := `{{.SnippetCore}} 
@@ -89,79 +96,82 @@ func GetHARHeadersCookies(headers *gabs.Container) (*gabs.Container, *gabs.Conta
 	return headersData, cookiesData
 }
 
-func GetRequestHARString(parsedAPI *gabs.Container) string {
+func GetRequestHARString(block *gabs.Container) string {
+	preprocess.ProcessVarsInBlock(block, globalVm)
+	httpv := block.S("verb", "value")
+	url := block.S("url", "value")
+	jsonObj := block.S("details", "ip_data")
+	headers := block.S("details", "headers")
+	/*
+		TODO: Handle multipart case
+
+		multipart := block.S("multipart", "value")
+		multipartBool := false
+		if multipart != nil {
+			multipartBool = true
+		}
+	*/
+	harObj := gabs.New()
+
+	if jsonObj != nil {
+		postData := gabs.New()
+		postData.Set("application/json", "mimeType")
+		postData.Set(jsonObj.String(), "text")
+		harObj.Set(postData, "postData")
+	}
+
+	if headers != nil {
+		headersData, cookiesData := GetHARHeadersCookies(headers)
+		if cookiesData.String() != "[]" {
+			harObj.Set(cookiesData, "cookies")
+		}
+		harObj.Set(headersData, "headers")
+	}
+
+	harObj.Set(httpv, "method")
+	harObj.Set(url, "url")
+
+	res := harObj.String()
+	return res
+}
+
+func GenerateTargetCode(targetLangLib string, parsedAPI *gabs.Container) {
+	initialize()
 	parsedAPIblocks := parsedAPI.S("value").Data().(*gabs.Container).Children()
-	vm := cmdexec.GetJSVm()
+	convertedSnippetList := make([]string, 0)
+
 	for i, block := range parsedAPIblocks {
 		log.Debug().Int("Block num", i).Msg("")
 		log.Debug().Str("Block getting processed", block.String()).Msg("")
 		blockType := block.S("type").Data().(string)
 		if blockType == "processor" {
-			fmt.Println("Skipping processor block")
+			snippet := block.S("value").Data().(*gabs.Container).Data().(string)
+			convertedSnippetList = append(convertedSnippetList, snippet)
 		} else if blockType == "Lama2File" {
-			preprocess.ProcessVarsInBlock(block, vm)
-			httpv := block.S("verb", "value")
-			url := block.S("url", "value")
-			jsonObj := block.S("details", "ip_data")
-			headers := block.S("details", "headers")
-			/*
-				TODO: Handle multipart case
-
-				multipart := block.S("multipart", "value")
-				multipartBool := false
-				if multipart != nil {
-					multipartBool = true
-				}
-			*/
-			harObj := gabs.New()
-
-			if jsonObj != nil {
-				postData := gabs.New()
-				postData.Set("application/json", "mimeType")
-				postData.Set(jsonObj.String(), "text")
-				harObj.Set(postData, "postData")
+			harRequest := GetRequestHARString(block)
+			snippetArgs := SnippetArgs{}
+			lang, lib := SplitLangLib(targetLangLib)
+			snippetArgs.Language = lang
+			snippetArgs.Library = lib
+			snippetArgs.HARRequest = harRequest
+			snippetArgs.SnippetCore = snippetcore
+			httpsnippetCode := PrepareHTTPSnippetGenerator(snippetArgs)
+
+			vm := cmdexec.GetJSVm()
+			_, e := vm.RunString(httpsnippetCode)
+			if e != nil {
+				log.Fatal().
+					Str("Type", "CodeGen").
+					Str("Error", e.Error()).
+					Msg("Code generator error")
 			}
-
-			if headers != nil {
-				headersData, cookiesData := GetHARHeadersCookies(headers)
-				if cookiesData.String() != "[]" {
-					harObj.Set(cookiesData, "cookies")
-				}
-				harObj.Set(headersData, "headers")
-			}
-
-			harObj.Set(httpv, "method")
-			harObj.Set(url, "url")
-
-			res := harObj.String()
-			return res
+			// Init returns an error if the package is not ready for use.
+			convertedSnippet := vm.Get("convertedSnippet").String()
+			convertedSnippetList = append(convertedSnippetList, convertedSnippet)
 		}
 	}
-	// TODO: Handle multi-stage files properly
-	// Better exception handling
-	return ""
-}
-
-func GenerateTargetCode(targetLangLib string, parsedAPI *gabs.Container) {
-	harRequest := GetRequestHARString(parsedAPI)
-	snippetArgs := SnippetArgs{}
-	lang, lib := SplitLangLib(targetLangLib)
-	snippetArgs.Language = lang
-	snippetArgs.Library = lib
-	snippetArgs.HARRequest = harRequest
-	snippetArgs.SnippetCore = snippetcore
-	httpsnippetCode := PrepareHTTPSnippetGenerator(snippetArgs)
-	vm := cmdexec.GetJSVm()
-	_, e := vm.RunString(httpsnippetCode)
-	if e != nil {
-		log.Fatal().
-			Str("Type", "CodeGen").
-			Str("Error", e.Error()).
-			Msg(fmt.Sprint("Code generator error"))
-	}
-	// Init returns an error if the package is not ready for use.
-	convertedSnippet := vm.Get("convertedSnippet").String()
-	fmt.Println(convertedSnippet)
-	clipboard.WriteAll(convertedSnippet)
+	convertedSnippetFinal := strings.Join(convertedSnippetList, "\n")
+	fmt.Println(convertedSnippetFinal)
+	clipboard.WriteAll(convertedSnippetFinal)
 	fmt.Println("Code copied to clipboard")
 }

controller/controller.go

@@ -99,7 +99,7 @@ func Process(version string) {
 	}
 
 	preprocess.LoadEnvironments(dir)
-	utils.ChangeWorkingDir(oldDir)
+	defer utils.ChangeWorkingDir(oldDir)
 	p := parser.NewLama2Parser()
 	parsedAPI, e := p.Parse(apiContent)
 	if o.Convert != "" {
@@ -117,7 +117,7 @@ func Process(version string) {
 			Str("Type", "Controller").
 			Str("LamaFile", o.Positional.LamaAPIFile).
 			Str("Error", e.Error()).
-			Msg(fmt.Sprint("Parse Error"))
+			Msg("Parse Error")
 	}
 	log.Debug().Str("Parsed API", parsedAPI.String()).Msg("")
 	HandleParsedFile(parsedAPI, o, dir)

tests/codegen_test.go

@@ -1,9 +1,14 @@
 package tests
 
 import (
+	"fmt"
+	"os"
 	"testing"
 
 	"github.com/HexmosTech/lama2/codegen"
+	"github.com/HexmosTech/lama2/parser"
+	"github.com/HexmosTech/lama2/preprocess"
+	"github.com/HexmosTech/lama2/utils"
 )
 
 func TestSplitLangLib1(t *testing.T) {
@@ -27,3 +32,18 @@ func TestSplitLangLib2(t *testing.T) {
 		t.Errorf("Expected lib = ''")
 	}
 }
+
+func TestGenerateMultiStage(t *testing.T) {
+	l2Path := "../examples/0009_processor_basic/0009_processor_basic.l2"
+	apiContent, _ := os.ReadFile(l2Path)
+	_, dir, _ := utils.GetFilePathComponents(l2Path)
+	nowPwd, _ := os.Getwd()
+	utils.ChangeWorkingDir(dir)
+	defer utils.ChangeWorkingDir(nowPwd)
+	preprocess.LoadEnvironments(dir)
+	p := parser.NewLama2Parser()
+	parsedAPI, _ := p.Parse(string(apiContent))
+	fmt.Println(parsedAPI)
+
+	codegen.GenerateTargetCode("python.requests", parsedAPI)
+}