Add flag to not add top-level modules to PYTHONPATH

stub_template_host.txt added all the top-level modules to the
PYTHONPATH, which isn't correct, and caused absl.logging to
override the built-in logging module.

Removing this also makes it more consistent with python binaries
built with embedded_launcher: true. embedded_launcher: true
binaries don't add the top-level modules.

Fixes: 245583294
Test: m py_dont_add_top_level_dirs_test && out/host/linux-x86/testcases/py_dont_add_top_level_dirs_test/x86_64/py_dont_add_top_level_dirs_test
Change-Id: Id3069565d2b2c4b2bda0ff5301e757a7b4201751
diff --git a/python/binary.go b/python/binary.go
index af29bb6..9c8c1f4 100644
--- a/python/binary.go
+++ b/python/binary.go
@@ -116,6 +116,14 @@
 	// doesn't exist next to the Android.bp, this attribute doesn't need to be set to true
 	// explicitly.
 	Auto_gen_config *bool
+
+	// Currently, both the root of the zipfile and all the directories 1 level
+	// below that are added to the python path. When this flag is set to true,
+	// only the root of the zipfile will be added to the python path. This flag
+	// will be removed after all the python modules in the tree have been updated
+	// 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
 }
 
 type binaryDecorator struct {
@@ -128,10 +136,6 @@
 	IntermPathForModuleOut() android.OptionalPath
 }
 
-var (
-	StubTemplateHost = "build/soong/python/scripts/stub_template_host.txt"
-)
-
 func NewBinary(hod android.HostOrDeviceSupported) (*Module, *binaryDecorator) {
 	module := newModule(hod, android.MultilibFirst)
 	decorator := &binaryDecorator{pythonInstaller: NewPythonInstaller("bin", "")}
@@ -180,9 +184,12 @@
 		})
 	}
 
+	addTopDirectoriesToPath := !proptools.BoolDefault(binary.binaryProperties.Dont_add_top_level_directories_to_path, false)
+
 	binFile := registerBuildActionForParFile(ctx, embeddedLauncher, launcherPath,
 		binary.getHostInterpreterName(ctx, actualVersion),
-		main, binary.getStem(ctx), append(android.Paths{srcsZip}, depsSrcsZips...))
+		main, binary.getStem(ctx), append(android.Paths{srcsZip}, depsSrcsZips...),
+		addTopDirectoriesToPath)
 
 	return android.OptionalPathForPath(binFile)
 }
diff --git a/python/builder.go b/python/builder.go
index 7d7239c..08d345c 100644
--- a/python/builder.go
+++ b/python/builder.go
@@ -44,13 +44,13 @@
 
 	hostPar = pctx.AndroidStaticRule("hostPar",
 		blueprint.RuleParams{
-			Command: `sed -e 's/%interpreter%/$interp/g' -e 's/%main%/$main/g' $template > $stub && ` +
+			Command: `sed -e 's/%interpreter%/$interp/g' -e 's/%main%/$main/g' -e 's/ADD_TOP_DIRECTORIES_TO_PATH/$addTopDirectoriesToPath/g' build/soong/python/scripts/stub_template_host.txt > $out.main && ` +
 				`echo "#!/usr/bin/env $interp" >${out}.prefix &&` +
-				`$mergeParCmd -p --prefix ${out}.prefix -pm $stub $out $srcsZips && ` +
-				`chmod +x $out && (rm -f $stub; rm -f ${out}.prefix)`,
-			CommandDeps: []string{"$mergeParCmd"},
+				`$mergeParCmd -p --prefix ${out}.prefix -pm $out.main $out $srcsZips && ` +
+				`chmod +x $out && (rm -f $out.main; rm -f ${out}.prefix)`,
+			CommandDeps: []string{"$mergeParCmd", "build/soong/python/scripts/stub_template_host.txt"},
 		},
-		"interp", "main", "template", "stub", "srcsZips")
+		"interp", "main", "srcsZips", "addTopDirectoriesToPath")
 
 	embeddedPar = pctx.AndroidStaticRule("embeddedPar",
 		blueprint.RuleParams{
@@ -81,7 +81,7 @@
 
 func registerBuildActionForParFile(ctx android.ModuleContext, embeddedLauncher bool,
 	launcherPath android.OptionalPath, interpreter, main, binName string,
-	srcsZips android.Paths) android.Path {
+	srcsZips android.Paths, addTopDirectoriesToPath bool) android.Path {
 
 	// .intermediate output path for bin executable.
 	binFile := android.PathForModuleOut(ctx, binName)
@@ -90,24 +90,20 @@
 	implicits := srcsZips
 
 	if !embeddedLauncher {
-		// the path of stub_template_host.txt from source tree.
-		template := android.PathForSource(ctx, StubTemplateHost)
-		implicits = append(implicits, template)
-
-		// intermediate output path for __main__.py
-		stub := android.PathForModuleOut(ctx, mainFileName).String()
-
+		addDirsString := "False"
+		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),
-				"template": template.String(),
-				"stub":     stub,
-				"srcsZips": strings.Join(srcsZips.Strings(), " "),
+				"interp":                  strings.Replace(interpreter, "/", `\/`, -1),
+				"main":                    strings.Replace(main, "/", `\/`, -1),
+				"srcsZips":                strings.Join(srcsZips.Strings(), " "),
+				"addTopDirectoriesToPath": addDirsString,
 			},
 		})
 	} else if launcherPath.Valid() {
diff --git a/python/python.go b/python/python.go
index daf7c14..f6029c2 100644
--- a/python/python.go
+++ b/python/python.go
@@ -356,10 +356,6 @@
 	protoExt             = ".proto"
 	pyVersion2           = "PY2"
 	pyVersion3           = "PY3"
-	initFileName         = "__init__.py"
-	mainFileName         = "__main__.py"
-	entryPointFile       = "entry_point.txt"
-	parFileExt           = ".zip"
 	internalPath         = "internal"
 )
 
diff --git a/python/python_test.go b/python/python_test.go
index f57f504..42a1ffb 100644
--- a/python/python_test.go
+++ b/python/python_test.go
@@ -300,8 +300,6 @@
 				filepath.Join("dir", "file2.py"):       nil,
 				filepath.Join("dir", "bin.py"):         nil,
 				filepath.Join("dir", "file4.py"):       nil,
-				StubTemplateHost: []byte(`PYTHON_BINARY = '%interpreter%'
-				MAIN_FILE = '%main%'`),
 			},
 			expectedBinaries: []pyModule{
 				{
diff --git a/python/scripts/stub_template_host.txt b/python/scripts/stub_template_host.txt
index 23897b3..a0ddffe 100644
--- a/python/scripts/stub_template_host.txt
+++ b/python/scripts/stub_template_host.txt
@@ -1,7 +1,6 @@
 #!/usr/bin/env '%interpreter%'
 
 import os
-import re
 import tempfile
 import shutil
 import sys
@@ -15,56 +14,31 @@
 # Don't imply 'import site' on initialization
 PYTHON_ARG = '-S'
 
-def SearchPathEnv(name):
-  search_path = os.getenv('PATH', os.defpath).split(os.pathsep)
-  for directory in search_path:
-    if directory == '': continue
-    path = os.path.join(directory, name)
-    # Check if path is actual executable file.
-    if os.path.isfile(path) and os.access(path, os.X_OK):
-      return path
-  return None
-
-def FindPythonBinary():
-  if PYTHON_BINARY.startswith('/'):
-    # Case 1: Python interpreter is directly provided with absolute path.
-    return PYTHON_BINARY
-  else:
-    # Case 2: Find Python interpreter through environment variable: PATH.
-    return SearchPathEnv(PYTHON_BINARY)
-
-# Create the runfiles tree by extracting the zip file
-def ExtractRunfiles():
-  temp_dir = tempfile.mkdtemp("", "Soong.python_")
-  zf = zipfile.ZipFile(os.path.dirname(__file__))
-  zf.extractall(temp_dir)
-  return temp_dir
-
 def Main():
   args = sys.argv[1:]
 
-  new_env = {}
-  runfiles_path = None
-
+  runfiles_path = tempfile.mkdtemp(prefix="Soong.python_")
   try:
-    runfiles_path = ExtractRunfiles()
+    zf = zipfile.ZipFile(os.path.dirname(__file__))
+    zf.extractall(runfiles_path)
+    zf.close()
 
     # Add runfiles path to PYTHONPATH.
     python_path_entries = [runfiles_path]
 
-    # Add top dirs within runfiles path to PYTHONPATH.
-    top_entries = [os.path.join(runfiles_path, i) for i in os.listdir(runfiles_path)]
-    top_pkg_dirs = [i for i in top_entries if os.path.isdir(i)]
-    python_path_entries += top_pkg_dirs
+    if ADD_TOP_DIRECTORIES_TO_PATH:
+      # Add top dirs within runfiles path to PYTHONPATH.
+      top_entries = [os.path.join(runfiles_path, i) for i in os.listdir(runfiles_path)]
+      top_pkg_dirs = [i for i in top_entries if os.path.isdir(i)]
+      python_path_entries += top_pkg_dirs
 
+    new_python_path = ":".join(python_path_entries)
     old_python_path = os.environ.get(PYTHON_PATH)
-    separator = ':'
-    new_python_path = separator.join(python_path_entries)
 
-    # Copy old PYTHONPATH.
     if old_python_path:
-      new_python_path += separator + old_python_path
-    new_env[PYTHON_PATH] = new_python_path
+      os.environ.update({PYTHON_PATH: new_python_path + ":" + old_python_path})
+    else:
+      os.environ.update({PYTHON_PATH: new_python_path})
 
     # Now look for main python source file.
     main_filepath = os.path.join(runfiles_path, MAIN_FILE)
@@ -73,23 +47,14 @@
     assert os.access(main_filepath, os.R_OK), \
            'Cannot exec() %r: file not readable.' % main_filepath
 
-    python_program = FindPythonBinary()
-    if python_program is None:
-      raise AssertionError('Could not find python binary: ' + PYTHON_BINARY)
-    args = [python_program, PYTHON_ARG, main_filepath] + args
-
-    os.environ.update(new_env)
+    args = [PYTHON_BINARY, PYTHON_ARG, main_filepath] + args
 
     sys.stdout.flush()
     # close_fds=False so that you can run binaries with files provided on the command line:
     # my_python_app --file <(echo foo)
-    retCode = subprocess.call(args, close_fds=False)
-    sys.exit(retCode)
-  except:
-    raise
+    sys.exit(subprocess.call(args, close_fds=False))
   finally:
-    if runfiles_path is not None:
-      shutil.rmtree(runfiles_path, True)
+    shutil.rmtree(runfiles_path, ignore_errors=True)
 
 if __name__ == '__main__':
   Main()
diff --git a/python/tests/top_level_dirs/Android.bp b/python/tests/top_level_dirs/Android.bp
new file mode 100644
index 0000000..0b15ce7
--- /dev/null
+++ b/python/tests/top_level_dirs/Android.bp
@@ -0,0 +1,9 @@
+python_test_host {
+    name: "py_dont_add_top_level_dirs_test",
+    main: "main.py",
+    srcs: [
+        "main.py",
+        "mypkg/mymodule.py",
+    ],
+    dont_add_top_level_directories_to_path: true,
+}
diff --git a/python/tests/top_level_dirs/main.py b/python/tests/top_level_dirs/main.py
new file mode 100644
index 0000000..9f30bfa
--- /dev/null
+++ b/python/tests/top_level_dirs/main.py
@@ -0,0 +1,17 @@
+import unittest
+import sys
+
+print(sys.path, file=sys.stderr)
+
+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/top_level_dirs/mypkg/mymodule.py b/python/tests/top_level_dirs/mypkg/mymodule.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/python/tests/top_level_dirs/mypkg/mymodule.py