From 95f35563ef2f2c4c0bdce29639dae6756540b212 Mon Sep 17 00:00:00 2001
From: Federico Mon <federico.mon@datadoghq.com>
Date: Tue, 23 Jun 2026 15:35:37 +0000
Subject: [PATCH] Handle PEP 440 local prerelease versions

---
 internal/platform/python.go      | 12 ++++++------
 internal/platform/python_test.go |  5 ++++-
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/internal/platform/python.go b/internal/platform/python.go
index b8e65c1..cc06dfa 100644
--- a/internal/platform/python.go
+++ b/internal/platform/python.go
@@ -17,15 +17,15 @@ import (
 	"github.com/DataDog/ddtest/internal/version"
 )
 
-// pep440PreReleaseRe matches PEP 440 pre-release suffixes (a/b/rc + digits) embedded
-// directly in a version component, e.g. "0rc1" or "12b3". We normalize these to
-// semver-style by inserting a hyphen so the existing parser can handle them.
-var pep440PreReleaseRe = regexp.MustCompile(`^(\d+\.\d+(?:\.\d+)*)((?:a|b|rc)\d+)$`)
+// pep440PreReleaseRe matches PEP 440 pre-release suffixes (a/b/rc + digits)
+// embedded directly in a Python package version, preserving optional local
+// version metadata, e.g. "4.12.0rc1+gabc123".
+var pep440PreReleaseRe = regexp.MustCompile(`^(\d+\.\d+(?:\.\d+)*)(a|b|rc)(\d+)(\+.*)?$`)
 
 // normalizePyVersion converts PEP 440 version strings to semver-compatible ones.
-// "4.12.0rc1" → "4.12.0-rc1", "4.10.3" → "4.10.3" (unchanged).
+// "4.12.0rc1+gabc123" -> "4.12.0-rc1+gabc123"; "4.10.3" is unchanged.
 func normalizePyVersion(v string) string {
-	return pep440PreReleaseRe.ReplaceAllString(v, "$1-$2")
+	return pep440PreReleaseRe.ReplaceAllString(v, "$1-$2$3$4")
 }
 
 //go:embed scripts/python_env.py
diff --git a/internal/platform/python_test.go b/internal/platform/python_test.go
index dd00215..f5278a9 100644
--- a/internal/platform/python_test.go
+++ b/internal/platform/python_test.go
@@ -22,8 +22,11 @@ func TestPython_Name(t *testing.T) {
 func TestNormalizePyVersion(t *testing.T) {
 	cases := []struct{ in, want string }{
 		{"4.12.0rc1", "4.12.0-rc1"},
+		{"4.12.0rc1+g0e3e598", "4.12.0-rc1+g0e3e598"},
 		{"4.12.0b2", "4.12.0-b2"},
+		{"4.12.0b2+gabc123", "4.12.0-b2+gabc123"},
 		{"4.12.0a1", "4.12.0-a1"},
+		{"4.12.0a1+local.build", "4.12.0-a1+local.build"},
 		{"4.10.3", "4.10.3"},
 		{"1.2.3.4", "1.2.3.4"},
 	}
@@ -36,7 +39,7 @@ func TestNormalizePyVersion(t *testing.T) {
 
 func TestPython_SanityCheck_SuccessWithPreRelease(t *testing.T) {
 	mockExecutor := &mockCommandExecutor{
-		combinedOutput: []byte("4.12.0rc1\n"),
+		combinedOutput: []byte("4.12.0rc1+g0e3e598\n"),
 	}
 	python := NewPython()
 	python.executor = mockExecutor