Add support for launching modules with pydev

Author: briandealwis

python/helper-image/launcher/launcher.go

@@ -52,6 +52,7 @@ import (
 	"flag"
 	"fmt"
 	"io"
+	"io/ioutil"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -312,15 +313,19 @@ func (pc *pythonContext) updateCommandLine(ctx context.Context) error {
 		if pc.env["WRAPPER_VERBOSE"] != "" {
 			cmdline = append(cmdline, "--DEBUG")
 		}
-		if pc.debugMode == ModePydevdPycharm {
-			// From the pydevd source, PyCharm wants multiproc
-			cmdline = append(cmdline, "--multiproc")
-		}
 		if !pc.wait {
 			cmdline = append(cmdline, "--continue")
 		}
-		cmdline = append(cmdline, "--file") // --file is expected as last argument
-		cmdline = append(cmdline, pc.args[1:]...)
+
+		// --file is expected as last pydev argument, but it must be a file, and so launching with
+		// a module requires some special handling.
+		cmdline = append(cmdline, "--file")
+		file, args, err := handlePydevModule(pc.args[1:])
+		if err != nil {
+			return err
+		}
+		cmdline = append(cmdline, file)
+		cmdline = append(cmdline, args...)
 		if pc.wait {
 			logrus.Warn("pydevd does not support wait-for-client")
 		}
@@ -356,6 +361,47 @@ func determinePythonMajorMinor(ctx context.Context, launcherBin string, env env)
 	return
 }
 
+// handlePydevModule applies special pydevd handling for a python module.  When a module is
+// found, we write out a python script that uses runpy to invoke the module.
+func handlePydevModule(args []string) (string, []string, error) {
+	switch {
+	case len(args) == 0:
+		return "", nil, fmt.Errorf("no python command-line specified") // shouldn't happen
+	case !strings.HasPrefix(args[0], "-"):
+		// this is a file
+		return args[0], args[1:], nil
+	case !strings.HasPrefix(args[0], "-m"):
+		// this is some other command-line flag
+		return "", nil, fmt.Errorf("expected python module: %q", args)
+	}
+	module := args[0][2:]
+	remaining := args[1:]
+	if module == "" {
+		if len(args) == 1 {
+			return "", nil, fmt.Errorf("missing python module: %q", args)
+		}
+		module = args[1]
+		remaining = args[2:]
+	}
+
+	snippet := strings.ReplaceAll(`import sys
+import runpy
+runpy.run_module('{module}', run_name="__main__",alter_sys=True)
+`, `{module}`, module)
+
+	// write out the temp location as other locations may not be writable
+	d, err := ioutil.TempDir("", "pydevd*")
+	if err != nil {
+		return "", nil, err
+	}
+	// use a skaffold-specific file name to ensure no possibility of it matching a user import
+	f := filepath.Join(d, "skaffold_pydevd_launch.py")
+	if err := ioutil.WriteFile(f, []byte(snippet), 0755); err != nil {
+		return "", nil, err
+	}
+	return f, remaining, nil
+}
+
 func isEnabled(env env) bool {
 	v, found := env["WRAPPER_ENABLED"]
 	return !found || (v != "0" && v != "false" && v != "no")

python/helper-image/launcher/launcher_test.go

@@ -20,10 +20,12 @@ import (
 	"context"
 	"fmt"
 	"io/ioutil"
+	"os"
 	"path/filepath"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
 )
 
 func TestValidateDebugMode(t *testing.T) {
@@ -312,3 +314,85 @@ func TestPathExists(t *testing.T) {
 		t.Error("pathExists failed on real path")
 	}
 }
+
+func TestHandlePydevModule(t *testing.T) {
+	tmp := os.TempDir()
+
+	tests := []struct {
+		description string
+		args        []string
+		shouldErr   bool
+		module      string
+		file        string
+		remaining   []string
+	}{
+		{
+			description: "plain file",
+			args:        []string{"app.py"},
+			file:        "app.py",
+		},
+		{
+			description: "-mmodule",
+			args:        []string{"-mmodule"},
+			file:        filepath.Join(tmp, "*", "skaffold_pydevd_launch.py"),
+		},
+		{
+			description: "-m module",
+			args:        []string{"-m", "module"},
+			file:        filepath.Join(tmp, "*", "skaffold_pydevd_launch.py"),
+		},
+		{
+			description: "- should error",
+			args:        []string{"-", "module"},
+			shouldErr:   true,
+		},
+		{
+			description: "-x should error",
+			args:        []string{"-x", "module"},
+			shouldErr:   true,
+		},
+		{
+			description: "lone -m should error",
+			args:        []string{"-m"},
+			shouldErr:   true,
+		},
+		{
+			description: "no args should error",
+			shouldErr:   true,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.description, func(t *testing.T) {
+			file, args, err := handlePydevModule(test.args)
+			if test.shouldErr {
+				if err == nil {
+					t.Error("Expected an error")
+				}
+			} else {
+				if !fileMatch(t, test.file, file) {
+					t.Errorf("Wanted %q but got %q", test.file, file)
+				}
+				if diff := cmp.Diff(args, test.remaining, cmpopts.EquateEmpty()); diff != "" {
+					t.Errorf("remaining args %T differ (-got, +want): %s", test.remaining, diff)
+				}
+			}
+		})
+	}
+}
+
+func fileMatch(t *testing.T, glob, file string) bool {
+	if file == glob {
+		return true
+	}
+	matches, err := filepath.Glob(glob)
+	if err != nil {
+		t.Errorf("Failed to expand globe %q: %v", glob, err)
+		return false
+	}
+	for _, m := range matches {
+		if file == m {
+			return true
+		}
+	}
+	return false
+}