Have python_*{,_host} handle arch-variants

Bug: 196081778
Test: TestPython*{,Host}ArchVariance
Test: go test
Test: mixed_{libc,droid}.sh
Change-Id: I89304e58f5bacd61534732bade4ad6bb5f2671c0

5 files changed, 126 insertions, 14 deletions

diff --git a/bp2build/python_binary_conversion_test.go b/bp2build/python_binary_conversion_test.go
index 6f6fc11e5..5b4829eb0 100644
--- a/bp2build/python_binary_conversion_test.go
+++ b/bp2build/python_binary_conversion_test.go
@@ -116,3 +116,37 @@ func TestPythonBinaryHostPy3(t *testing.T) {
 		},
 	})
 }
+
+func TestPythonBinaryHostArchVariance(t *testing.T) {
+	runBp2BuildTestCaseSimple(t, bp2buildTestCase{
+		description:                        "test arch variants",
+		moduleTypeUnderTest:                "python_binary_host",
+		moduleTypeUnderTestFactory:         python.PythonBinaryHostFactory,
+		moduleTypeUnderTestBp2BuildMutator: python.PythonBinaryBp2Build,
+		filesystem: map[string]string{
+			"dir/arm.py": "",
+			"dir/x86.py": "",
+		},
+		blueprint: `python_binary_host {
+					 name: "foo-arm",
+					 arch: {
+						 arm: {
+							 srcs: ["arm.py"],
+						 },
+						 x86: {
+							 srcs: ["x86.py"],
+						 },
+					},
+				 }`,
+		expectedBazelTargets: []string{
+			`py_binary(
+    name = "foo-arm",
+    srcs = select({
+        "//build/bazel/platforms/arch:arm": ["arm.py"],
+        "//build/bazel/platforms/arch:x86": ["x86.py"],
+        "//conditions:default": [],
+    }),
+)`,
+		},
+	})
+}
diff --git a/bp2build/python_library_conversion_test.go b/bp2build/python_library_conversion_test.go
index b6f45e53e..7f983ad50 100644
--- a/bp2build/python_library_conversion_test.go
+++ b/bp2build/python_library_conversion_test.go
@@ -154,3 +154,54 @@ func testPythonLib(t *testing.T, modType string,
 		},
 	})
 }
+
+func TestPythonLibraryArchVariance(t *testing.T) {
+	testPythonArchVariance(t, "python_library", "py_library",
+		python.PythonLibraryFactory, python.PythonLibraryBp2Build,
+		func(ctx android.RegistrationContext) {})
+}
+
+func TestPythonLibraryHostArchVariance(t *testing.T) {
+	testPythonArchVariance(t, "python_library_host", "py_library",
+		python.PythonLibraryHostFactory, python.PythonLibraryHostBp2Build,
+		func(ctx android.RegistrationContext) {})
+}
+
+// TODO: refactor python_binary_conversion_test to use this
+func testPythonArchVariance(t *testing.T, modType, bazelTarget string,
+	factory android.ModuleFactory, mutator PythonLibBp2Build,
+	registration func(ctx android.RegistrationContext)) {
+	t.Helper()
+	runBp2BuildTestCase(t, registration, bp2buildTestCase{
+		description:                        fmt.Sprintf("test %s arch variants", modType),
+		moduleTypeUnderTest:                modType,
+		moduleTypeUnderTestFactory:         factory,
+		moduleTypeUnderTestBp2BuildMutator: mutator,
+		filesystem: map[string]string{
+			"dir/arm.py": "",
+			"dir/x86.py": "",
+		},
+		blueprint: fmt.Sprintf(`%s {
+					 name: "foo",
+					 arch: {
+						 arm: {
+							 srcs: ["arm.py"],
+						 },
+						 x86: {
+							 srcs: ["x86.py"],
+						 },
+					},
+				 }`, modType),
+		expectedBazelTargets: []string{
+			fmt.Sprintf(`%s(
+    name = "foo",
+    srcs = select({
+        "//build/bazel/platforms/arch:arm": ["arm.py"],
+        "//build/bazel/platforms/arch:x86": ["x86.py"],
+        "//conditions:default": [],
+    }),
+    srcs_version = "PY3",
+)`, bazelTarget),
+		},
+	})
+}
diff --git a/python/binary.go b/python/binary.go
index bc2768c72..afcc53a65 100644
--- a/python/binary.go
+++ b/python/binary.go
@@ -63,6 +63,7 @@ func PythonBinaryBp2Build(ctx android.TopDownMutatorContext) {
 			}
 		}
 	}
+
 	// TODO(b/182306917): this doesn't fully handle all nested props versioned
 	// by the python version, which would have been handled by the version split
 	// mutator. This is sufficient for very simple python_binary_host modules
@@ -80,15 +81,12 @@ func PythonBinaryBp2Build(ctx android.TopDownMutatorContext) {
 		// do nothing, since python_version defaults to PY3.
 	}
 
-	srcs := android.BazelLabelForModuleSrcExcludes(ctx, m.properties.Srcs, m.properties.Exclude_srcs)
-	data := android.BazelLabelForModuleSrc(ctx, m.properties.Data)
-	deps := android.BazelLabelForModuleDeps(ctx, m.properties.Libs)
-
+	baseAttrs := m.makeArchVariantBaseAttributes(ctx)
 	attrs := &bazelPythonBinaryAttributes{
 		Main:           main,
-		Srcs:           bazel.MakeLabelListAttribute(srcs),
-		Data:           bazel.MakeLabelListAttribute(data),
-		Deps:           bazel.MakeLabelListAttribute(deps),
+		Srcs:           baseAttrs.Srcs,
+		Data:           baseAttrs.Data,
+		Deps:           baseAttrs.Deps,
 		Python_version: python_version,
 	}
 
diff --git a/python/library.go b/python/library.go
index a132216f8..19fa59a3f 100644
--- a/python/library.go
+++ b/python/library.go
@@ -88,14 +88,11 @@ func pythonLibBp2Build(ctx android.TopDownMutatorContext, modType string) {
 		// do nothing, since python_version defaults to PY2ANDPY3
 	}
 
-	srcs := android.BazelLabelForModuleSrcExcludes(ctx, m.properties.Srcs, m.properties.Exclude_srcs)
-	data := android.BazelLabelForModuleSrc(ctx, m.properties.Data)
-	deps := android.BazelLabelForModuleDeps(ctx, m.properties.Libs)
-
+	baseAttrs := m.makeArchVariantBaseAttributes(ctx)
 	attrs := &bazelPythonLibraryAttributes{
-		Srcs:         bazel.MakeLabelListAttribute(srcs),
-		Data:         bazel.MakeLabelListAttribute(data),
-		Deps:         bazel.MakeLabelListAttribute(deps),
+		Srcs:         baseAttrs.Srcs,
+		Data:         baseAttrs.Data,
+		Deps:         baseAttrs.Deps,
 		Srcs_version: python_version,
 	}
 
diff --git a/python/python.go b/python/python.go
index f90017268..401d91fe3 100644
--- a/python/python.go
+++ b/python/python.go
@@ -22,6 +22,7 @@ import (
 	"regexp"
 	"strings"
 
+	"android/soong/bazel"
 	"github.com/google/blueprint"
 	"github.com/google/blueprint/proptools"
 
@@ -120,6 +121,18 @@ type BaseProperties struct {
 	Embedded_launcher *bool `blueprint:"mutated"`
 }
 
+type baseAttributes struct {
+	// TODO(b/200311466): Probably not translate b/c Bazel has no good equiv
+	//Pkg_path    bazel.StringAttribute
+	// TODO: Related to Pkg_bath and similarLy gated
+	//Is_internal bazel.BoolAttribute
+	// Combines Srcs and Exclude_srcs
+	Srcs bazel.LabelListAttribute
+	Deps bazel.LabelListAttribute
+	// Combines Data and Java_data (invariant)
+	Data bazel.LabelListAttribute
+}
+
 // Used to store files of current module after expanding dependencies
 type pathMapping struct {
 	dest string
@@ -177,6 +190,25 @@ func newModule(hod android.HostOrDeviceSupported, multilib android.Multilib) *Mo
 	}
 }
 
+func (m *Module) makeArchVariantBaseAttributes(ctx android.TopDownMutatorContext) baseAttributes {
+	var attrs baseAttributes
+	archVariantBaseProps := m.GetArchVariantProperties(ctx, &BaseProperties{})
+	for axis, configToProps := range archVariantBaseProps {
+		for config, props := range configToProps {
+			if baseProps, ok := props.(*BaseProperties); ok {
+				attrs.Srcs.SetSelectValue(axis, config,
+					android.BazelLabelForModuleSrcExcludes(ctx, baseProps.Srcs, baseProps.Exclude_srcs))
+				attrs.Deps.SetSelectValue(axis, config,
+					android.BazelLabelForModuleDeps(ctx, baseProps.Libs))
+				data := android.BazelLabelForModuleSrc(ctx, baseProps.Data)
+				data.Append(android.BazelLabelForModuleSrc(ctx, baseProps.Java_data))
+				attrs.Data.SetSelectValue(axis, config, data)
+			}
+		}
+	}
+	return attrs
+}
+
 // bootstrapper interface should be implemented for runnable modules, e.g. binary and test
 type bootstrapper interface {
 	bootstrapperProps() []interface{}