From 476439c48ef14ceb433d399687080582b022775c Mon Sep 17 00:00:00 2001
From: Andrew Nester <andrew.nester@databricks.com>
Date: Thu, 17 Aug 2023 17:50:00 +0200
Subject: [PATCH] moved to separate function

---
 bundle/python/transform.go | 59 ++++++++++++++++++++++----------------
 1 file changed, 34 insertions(+), 25 deletions(-)

diff --git a/bundle/python/transform.go b/bundle/python/transform.go
index aa1a3db5..cd7c826f 100644
--- a/bundle/python/transform.go
+++ b/bundle/python/transform.go
@@ -53,38 +53,47 @@ func (m *transform) Name() string {
 func (m *transform) Apply(ctx context.Context, b *bundle.Bundle) error {
 	wheelTasks := libraries.FindAllWheelTasks(b)
 	for _, wheelTask := range wheelTasks {
-		taskDefinition := wheelTask.PythonWheelTask
-		libraries := wheelTask.Libraries
-
-		wheelTask.PythonWheelTask = nil
-		wheelTask.Libraries = nil
-
-		filename, err := generateNotebookWrapper(b, taskDefinition, libraries)
+		err := generateNotebookTrampoline(b, wheelTask)
 		if err != nil {
 			return err
 		}
-
-		internalDir, err := getInternalDir(b)
-		if err != nil {
-			return err
-		}
-
-		internalDirRel, err := filepath.Rel(b.Config.Path, internalDir)
-		if err != nil {
-			return err
-		}
-
-		parts := []string{b.Config.Workspace.FilesPath}
-		parts = append(parts, strings.Split(internalDirRel, string(os.PathSeparator))...)
-		parts = append(parts, filename)
-
-		wheelTask.NotebookTask = &jobs.NotebookTask{
-			NotebookPath: path.Join(parts...),
-		}
 	}
 	return nil
 }
 
+func generateNotebookTrampoline(b *bundle.Bundle, wheelTask *jobs.Task) error {
+	taskDefinition := wheelTask.PythonWheelTask
+	libraries := wheelTask.Libraries
+
+	wheelTask.PythonWheelTask = nil
+	wheelTask.Libraries = nil
+
+	filename, err := generateNotebookWrapper(b, taskDefinition, libraries)
+	if err != nil {
+		return err
+	}
+
+	internalDir, err := getInternalDir(b)
+	if err != nil {
+		return err
+	}
+
+	internalDirRel, err := filepath.Rel(b.Config.Path, internalDir)
+	if err != nil {
+		return err
+	}
+
+	parts := []string{b.Config.Workspace.FilesPath}
+	parts = append(parts, strings.Split(internalDirRel, string(os.PathSeparator))...)
+	parts = append(parts, filename)
+
+	wheelTask.NotebookTask = &jobs.NotebookTask{
+		NotebookPath: path.Join(parts...),
+	}
+
+	return nil
+}
+
 func getInternalDir(b *bundle.Bundle) (string, error) {
 	cacheDir, err := b.CacheDir()
 	if err != nil {