Add flag to not add directory of entrypoint to sys.path

The python interpreter will by default add the directory of
the entrypoint script to the beginning of sys.path. This
can be disabled in python 3.11+ (which is not released yet)
using the PYTHON_SAFE_PATH environment variable or the -P flag.
As a workaround to have this behavior in older python versions,
we can make an __soong_entrypoint_redirector__.py file at the
root of the zip file that is the entrypoint, and then that
file will redirect to the real entrypoint.

This brings non-embedded-launcher python modules closer to
the embedded launcher version. The embedded launcher binaries
already act like this because they start at an __main__.py file
at the root of the zip file.

Bug: 245583294
Test: m py_dont_import_folder_of_entrypoint_test && out/host/linux-x86/nativetest64/py_dont_import_folder_of_entrypoint_test/py_dont_import_folder_of_entrypoint_test
Change-Id: I39aaf04fb19c3ba7f5c9d98220872d6d08abf736

6 files changed, 96 insertions, 16 deletions

diff --git a/python/binary.go b/python/binary.go
index 9c8c1f43b..e6324a3b5 100644
--- a/python/binary.go
+++ b/python/binary.go
@@ -124,6 +124,14 @@ type BinaryProperties struct {
 	// to support it. When using embedded_launcher: true, this is already the
 	// behavior. The default is currently false.
 	Dont_add_top_level_directories_to_path *bool
+
+	// Setting this to true will mimic Python 3.11+'s PYTHON_SAFE_PATH environment
+	// variable or -P flag, even on older python versions. This is a temporary
+	// flag while modules are changed to support it, eventually true will be the
+	// default and the flag will be removed. The default is currently false. It
+	// is only applicable when embedded_launcher is false, when embedded_launcher
+	// is true this is already implied.
+	Dont_add_entrypoint_folder_to_path *bool
 }
 
 type binaryDecorator struct {
@@ -185,11 +193,12 @@ func (binary *binaryDecorator) bootstrap(ctx android.ModuleContext, actualVersio
 	}
 
 	addTopDirectoriesToPath := !proptools.BoolDefault(binary.binaryProperties.Dont_add_top_level_directories_to_path, false)
+	dontAddEntrypointFolderToPath := proptools.BoolDefault(binary.binaryProperties.Dont_add_entrypoint_folder_to_path, false)
 
 	binFile := registerBuildActionForParFile(ctx, embeddedLauncher, launcherPath,
 		binary.getHostInterpreterName(ctx, actualVersion),
 		main, binary.getStem(ctx), append(android.Paths{srcsZip}, depsSrcsZips...),
-		addTopDirectoriesToPath)
+		addTopDirectoriesToPath, dontAddEntrypointFolderToPath)
 
 	return android.OptionalPathForPath(binFile)
 }
diff --git a/python/builder.go b/python/builder.go
index 08d345c92..f7f9a9914 100644
--- a/python/builder.go
+++ b/python/builder.go
@@ -20,7 +20,6 @@ import (
 	"strings"
 
 	"android/soong/android"
-
 	"github.com/google/blueprint"
 	_ "github.com/google/blueprint/bootstrap"
 )
@@ -52,13 +51,25 @@ var (
 		},
 		"interp", "main", "srcsZips", "addTopDirectoriesToPath")
 
+	hostParWithoutAddingEntrypointFolderToPath = pctx.AndroidStaticRule("hostParWithoutAddingEntrypointFolderToPath",
+		blueprint.RuleParams{
+			Command: `sed -e 's/%interpreter%/$interp/g' -e 's/%main%/__soong_entrypoint_redirector__.py/g' -e 's/ADD_TOP_DIRECTORIES_TO_PATH/$addTopDirectoriesToPath/g' build/soong/python/scripts/stub_template_host.txt > $out.main && ` +
+				"sed -e 's/ENTRY_POINT/$main/g' build/soong/python/scripts/main_non_embedded.py >`dirname $out`/__soong_entrypoint_redirector__.py && " +
+				"$parCmd -o $out.entrypoint_zip -C `dirname $out` -f `dirname $out`/__soong_entrypoint_redirector__.py && " +
+				`echo "#!/usr/bin/env $interp" >${out}.prefix &&` +
+				`$mergeParCmd -p --prefix ${out}.prefix -pm $out.main $out $srcsZips $out.entrypoint_zip && ` +
+				"chmod +x $out && (rm -f $out.main; rm -f ${out}.prefix; rm -f $out.entrypoint_zip; rm -f `dirname $out`/__soong_entrypoint_redirector__.py)",
+			CommandDeps: []string{"$mergeParCmd", "$parCmd", "build/soong/python/scripts/stub_template_host.txt", "build/soong/python/scripts/main_non_embedded.py"},
+		},
+		"interp", "main", "srcsZips", "addTopDirectoriesToPath")
+
 	embeddedPar = pctx.AndroidStaticRule("embeddedPar",
 		blueprint.RuleParams{
 			Command: `rm -f $out.main && ` +
 				`sed 's/ENTRY_POINT/$main/' build/soong/python/scripts/main.py >$out.main &&` +
 				`$mergeParCmd -p -pm $out.main --prefix $launcher $out $srcsZips && ` +
 				`chmod +x $out && rm -rf $out.main`,
-			CommandDeps: []string{"$mergeParCmd", "$parCmd", "build/soong/python/scripts/main.py"},
+			CommandDeps: []string{"$mergeParCmd", "build/soong/python/scripts/main.py"},
 		},
 		"main", "srcsZips", "launcher")
 
@@ -81,7 +92,7 @@ func init() {
 
 func registerBuildActionForParFile(ctx android.ModuleContext, embeddedLauncher bool,
 	launcherPath android.OptionalPath, interpreter, main, binName string,
-	srcsZips android.Paths, addTopDirectoriesToPath bool) android.Path {
+	srcsZips android.Paths, addTopDirectoriesToPath bool, dontAddEntrypointFolderToPath bool) android.Path {
 
 	// .intermediate output path for bin executable.
 	binFile := android.PathForModuleOut(ctx, binName)
@@ -94,18 +105,33 @@ func registerBuildActionForParFile(ctx android.ModuleContext, embeddedLauncher b
 		if addTopDirectoriesToPath {
 			addDirsString = "True"
 		}
-		ctx.Build(pctx, android.BuildParams{
-			Rule:        hostPar,
-			Description: "host python archive",
-			Output:      binFile,
-			Implicits:   implicits,
-			Args: map[string]string{
-				"interp":                  strings.Replace(interpreter, "/", `\/`, -1),
-				"main":                    strings.Replace(main, "/", `\/`, -1),
-				"srcsZips":                strings.Join(srcsZips.Strings(), " "),
-				"addTopDirectoriesToPath": addDirsString,
-			},
-		})
+		if dontAddEntrypointFolderToPath {
+			ctx.Build(pctx, android.BuildParams{
+				Rule:        hostParWithoutAddingEntrypointFolderToPath,
+				Description: "host python archive",
+				Output:      binFile,
+				Implicits:   implicits,
+				Args: map[string]string{
+					"interp":                  strings.Replace(interpreter, "/", `\/`, -1),
+					"main":                    strings.Replace(strings.TrimSuffix(main, pyExt), "/", ".", -1),
+					"srcsZips":                strings.Join(srcsZips.Strings(), " "),
+					"addTopDirectoriesToPath": addDirsString,
+				},
+			})
+		} else {
+			ctx.Build(pctx, android.BuildParams{
+				Rule:        hostPar,
+				Description: "host python archive",
+				Output:      binFile,
+				Implicits:   implicits,
+				Args: map[string]string{
+					"interp":                  strings.Replace(interpreter, "/", `\/`, -1),
+					"main":                    strings.Replace(main, "/", `\/`, -1),
+					"srcsZips":                strings.Join(srcsZips.Strings(), " "),
+					"addTopDirectoriesToPath": addDirsString,
+				},
+			})
+		}
 	} else if launcherPath.Valid() {
 		// added launcherPath to the implicits Ninja dependencies.
 		implicits = append(implicits, launcherPath.Path())
diff --git a/python/scripts/main_non_embedded.py b/python/scripts/main_non_embedded.py
new file mode 100644
index 000000000..ffbaaa8db
--- /dev/null
+++ b/python/scripts/main_non_embedded.py
@@ -0,0 +1,6 @@
+import runpy
+
+# The purpose of this file is to implement python 3.11+'s
+# PYTHON_SAFE_PATH / -P option on older python versions.
+
+runpy._run_module_as_main("ENTRY_POINT", alter_argv=False)
diff --git a/python/tests/dont_import_folder_of_entrypoint/Android.bp b/python/tests/dont_import_folder_of_entrypoint/Android.bp
new file mode 100644
index 000000000..40f522fc4
--- /dev/null
+++ b/python/tests/dont_import_folder_of_entrypoint/Android.bp
@@ -0,0 +1,24 @@
+python_test_host {
+    name: "py_dont_import_folder_of_entrypoint_test",
+    main: "mypkg/main.py",
+    srcs: [
+        "mypkg/main.py",
+        "mypkg/mymodule.py",
+    ],
+    dont_add_entrypoint_folder_to_path: true,
+    dont_add_top_level_directories_to_path: true,
+}
+
+python_test_host {
+    name: "py_dont_import_folder_of_entrypoint_test_embedded_launcher",
+    main: "mypkg/main.py",
+    srcs: [
+        "mypkg/main.py",
+        "mypkg/mymodule.py",
+    ],
+    version: {
+        py3: {
+            embedded_launcher: true,
+        },
+    },
+}
diff --git a/python/tests/dont_import_folder_of_entrypoint/mypkg/main.py b/python/tests/dont_import_folder_of_entrypoint/mypkg/main.py
new file mode 100644
index 000000000..c6a36edfb
--- /dev/null
+++ b/python/tests/dont_import_folder_of_entrypoint/mypkg/main.py
@@ -0,0 +1,15 @@
+import unittest
+import sys
+
+class TestProtoWithPkgPath(unittest.TestCase):
+
+    def test_cant_import_mymodule_directly(self):
+        with self.assertRaises(ImportError):
+            import mymodule
+
+    def test_can_import_mymodule_by_parent_package(self):
+        import mypkg.mymodule
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/python/tests/dont_import_folder_of_entrypoint/mypkg/mymodule.py b/python/tests/dont_import_folder_of_entrypoint/mypkg/mymodule.py
new file mode 100644
index 000000000..e69de29bb
--- /dev/null
+++ b/python/tests/dont_import_folder_of_entrypoint/mypkg/mymodule.py