diff --git a/python/binary.go b/python/binary.go
index e6324a3..1c2b555 100644
--- a/python/binary.go
+++ b/python/binary.go
@@ -34,7 +34,7 @@
 }
 
 type bazelPythonBinaryAttributes struct {
-	Main           *string
+	Main           *bazel.Label
 	Srcs           bazel.LabelListAttribute
 	Deps           bazel.LabelListAttribute
 	Python_version *string
@@ -42,17 +42,6 @@
 }
 
 func pythonBinaryBp2Build(ctx android.TopDownMutatorContext, m *Module) {
-	var main *string
-	for _, propIntf := range m.GetProperties() {
-		if props, ok := propIntf.(*BinaryProperties); ok {
-			// main is optional.
-			if props.Main != nil {
-				main = props.Main
-				break
-			}
-		}
-	}
-
 	// 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
@@ -72,13 +61,24 @@
 
 	baseAttrs := m.makeArchVariantBaseAttributes(ctx)
 	attrs := &bazelPythonBinaryAttributes{
-		Main:           main,
+		Main:           nil,
 		Srcs:           baseAttrs.Srcs,
 		Deps:           baseAttrs.Deps,
 		Python_version: python_version,
 		Imports:        baseAttrs.Imports,
 	}
 
+	for _, propIntf := range m.GetProperties() {
+		if props, ok := propIntf.(*BinaryProperties); ok {
+			// main is optional.
+			if props.Main != nil {
+				main := android.BazelLabelForModuleSrcSingle(ctx, *props.Main)
+				attrs.Main = &main
+				break
+			}
+		}
+	}
+
 	props := bazel.BazelTargetModuleProperties{
 		// Use the native py_binary rule.
 		Rule_class: "py_binary",
diff --git a/python/defaults.go b/python/defaults.go
index c54e7d0..3dc5bc4 100644
--- a/python/defaults.go
+++ b/python/defaults.go
@@ -19,7 +19,7 @@
 )
 
 func init() {
-	android.RegisterModuleType("python_defaults", defaultsFactory)
+	android.RegisterModuleType("python_defaults", DefaultsFactory)
 }
 
 type Defaults struct {
@@ -30,7 +30,7 @@
 func (d *Defaults) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 }
 
-func defaultsFactory() android.Module {
+func DefaultsFactory() android.Module {
 	module := &Defaults{}
 
 	module.AddProperties(