Added filebase64sha256 support for Python

pkg/codegen/python/gen_program.go

@@ -50,8 +50,11 @@ func GenerateProgram(program *pcl.Program) (map[string][]byte, hcl.Diagnostics,
 	// Linearize the nodes into an order appropriate for procedural code generation.
 	nodes := pcl.Linearize(program)
 
+	// Creating a list to store and later print helper methods if they turn out to be needed
+	preambleHelperMethods := codegen.NewStringSet()
+
 	var main bytes.Buffer
-	g.genPreamble(&main, program)
+	g.genPreamble(&main, program, preambleHelperMethods)
 	for _, n := range nodes {
 		g.genNode(&main, n)
 	}
@@ -112,7 +115,7 @@ func (g *generator) genComment(w io.Writer, comment syntax.Comment) {
 	}
 }
 
-func (g *generator) genPreamble(w io.Writer, program *pcl.Program) {
+func (g *generator) genPreamble(w io.Writer, program *pcl.Program, preambleHelperMethods codegen.StringSet) {
 	// Print the pulumi import at the top.
 	g.Fprintln(w, "import pulumi")
 
@@ -125,8 +128,13 @@ func (g *generator) genPreamble(w io.Writer, program *pcl.Program) {
 		}
 		diags := n.VisitExpressions(nil, func(n model.Expression) (model.Expression, hcl.Diagnostics) {
 			if call, ok := n.(*model.FunctionCallExpression); ok {
-				if i := g.getFunctionImports(call); i != "" {
-					importSet.Add(i)
+				if i := g.getFunctionImports(call); i[0] != "" {
+					for _, importPackage := range i {
+						importSet.Add(importPackage)
+					}
+				}
+				if helperMethodBody, ok := getHelperMethodIfNeeded(call.Name); ok {
+					preambleHelperMethods.Add(helperMethodBody)
 				}
 			}
 			return n, nil
@@ -152,6 +160,11 @@ func (g *generator) genPreamble(w io.Writer, program *pcl.Program) {
 		g.Fprintln(w, i)
 	}
 	g.Fprint(w, "\n")
+
+	// If we collected any helper methods that should be added, write them just before the main func
+	for _, preambleHelperMethodBody := range preambleHelperMethods.SortedValues() {
+		g.Fprintf(w, "%s\n\n", preambleHelperMethodBody)
+	}
 }
 
 func (g *generator) genNode(w io.Writer, n pcl.Node) {

pkg/codegen/python/gen_program_expressions.go

@@ -181,24 +181,25 @@ func functionName(tokenArg model.Expression) (string, string, string, hcl.Diagno
 	return makeValidIdentifier(pkg), strings.Replace(module, "/", ".", -1), title(member), diagnostics
 }
 
-var functionImports = map[string]string{
-	"fileArchive": "pulumi",
-	"fileAsset":   "pulumi",
-	"filebase64":  "base64",
-	"readDir":     "os",
-	"toBase64":    "base64",
-	"toJSON":      "json",
-	"sha1":        "hashlib",
+var functionImports = map[string][]string{
+	"fileArchive":      []string{"pulumi"},
+	"fileAsset":        []string{"pulumi"},
+	"filebase64":       []string{"base64"},
+	"filebase64sha256": []string{"base64", "hashlib"},
+	"readDir":          []string{"os"},
+	"toBase64":         []string{"base64"},
+	"toJSON":           []string{"json"},
+	"sha1":             []string{"hashlib"},
 }
 
-func (g *generator) getFunctionImports(x *model.FunctionCallExpression) string {
+func (g *generator) getFunctionImports(x *model.FunctionCallExpression) []string {
 	if x.Name != pcl.Invoke {
 		return functionImports[x.Name]
 	}
 
 	pkg, _, _, diags := functionName(x.Args[0])
 	contract.Assert(len(diags) == 0)
-	return "pulumi_" + pkg
+	return []string{"pulumi_" + pkg}
 }
 
 func (g *generator) GenFunctionCallExpression(w io.Writer, expr *model.FunctionCallExpression) {
@@ -222,6 +223,9 @@ func (g *generator) GenFunctionCallExpression(w io.Writer, expr *model.FunctionC
 		g.Fgenf(w, "pulumi.FileAsset(%.v)", expr.Args[0])
 	case "filebase64":
 		g.Fgenf(w, "(lambda path: base64.b64encode(open(path).read().encode()).decode())(%.v)", expr.Args[0])
+	case "filebase64sha256":
+		// Assuming the existence of the following helper method
+		g.Fgenf(w, "computeFilebase64sha256(%v)", expr.Args[0])
 	case pcl.Invoke:
 		pkg, module, fn, diags := functionName(expr.Args[0])
 		contract.Assert(len(diags) == 0)

pkg/codegen/python/gen_program_utils.go

@@ -0,0 +1,17 @@
+package python
+
+// Provides code for a method which will be placed in the program preamble if deemed
+// necessary. Because many tasks in Go such as reading a file require extensive error
+// handling, it is much prettier to encapsulate that error handling boilerplate as its
+// own function in the preamble.
+func getHelperMethodIfNeeded(functionName string) (string, bool) {
+	switch functionName {
+	case "filebase64sha256":
+		return `def computeFilebase64sha256(path):
+	fileData = open(path).read().encode()
+	hashedData = hashlib.sha256(fileData.encode()).digest()
+	return base64.b64encode(hashedData).decode()`, true
+	default:
+		return "", false
+	}
+}