Emit pyproject version ranges

rjcoulter22 · rjcoulter22 · commit a77c6870e674 · 2026-05-06T16:33:49.000-04:00
diff --git a/cmd/datadog-sbom-generator/__snapshots__/main_test.snap b/cmd/datadog-sbom-generator/__snapshots__/main_test.snap
@@ -2377,6 +2377,33 @@ Scanned <rootdir>/fixtures/integration-npm/with-workspace/yarn.lock file and fou
         ]
       }
     },
+    {
+      "bom-ref": "pkg:pypi/numpy",
+      "type": "library",
+      "name": "numpy",
+      "purl": "pkg:pypi/numpy",
+      "properties": [
+        {
+          "name": "datadog:is-direct",
+          "value": "true"
+        },
+        {
+          "name": "datadog:package-manager",
+          "value": "uv"
+        },
+        {
+          "name": "datadog:version-range",
+          "value": "/u003e=1.24"
+        }
+      ],
+      "evidence": {
+        "occurrences": [
+          {
+            "location": "{/"block/":{/"file_name/":/"pyproject.toml/",/"line_start/":8,/"line_end/":8,/"column_start/":5,/"column_end/":19,/"role/":/"manifest/"},/"name/":{/"file_name/":/"pyproject.toml/",/"line_start/":8,/"line_end/":8,/"column_start/":6,/"column_end/":11,/"role/":/"manifest/"},/"version/":{/"file_name/":/"pyproject.toml/",/"line_start/":8,/"line_end/":8,/"column_start/":11,/"column_end/":17,/"role/":/"manifest/"}}"
+          }
+        ]
+      }
+    },
     {
       "bom-ref": "pkg:pypi/pytest@8.2.0",
       "type": "library",
@@ -2436,7 +2463,7 @@ Scanned <rootdir>/fixtures/integration-npm/with-workspace/yarn.lock file and fou
 
 [TestRun/Scan_pyproject.toml_without_lock_file - 2]
 Scanning directory './fixtures/integration-pyproject', resolved absolute path '<rootdir>/fixtures/integration-pyproject'
-Scanned <rootdir>/fixtures/integration-pyproject/pyproject.toml file and found 3 packages
+Scanned <rootdir>/fixtures/integration-pyproject/pyproject.toml file and found 4 packages
 [reachability] Reachability analysis is disabled
 ---
 
diff --git a/pkg/lockfile/fixtures/pyproject-toml-extractor/poetry-pinned/pyproject.toml b/pkg/lockfile/fixtures/pyproject-toml-extractor/poetry-pinned/pyproject.toml
@@ -12,3 +12,4 @@ requests = "==2.28.0"
 flask = "==2.3.2"
 boto3 = "1.26.0"
 numpy = "^1.24"
+scipy = "1.*"
diff --git a/pkg/lockfile/python/parse-pyproject-toml.go b/pkg/lockfile/python/parse-pyproject-toml.go
@@ -3,6 +3,7 @@ package python
 import (
 	"fmt"
 	"io"
+	"log"
 	"maps"
 	"path/filepath"
 	"slices"
@@ -147,20 +148,25 @@ func (e PyProjectTOMLExtractor) Extract(f lockfile.DepFile, context lockfile.Sca
 
 	lines := fileposition.BytesToLines(content)
 	pm := detectPackageManager(&pyproject)
-	packages := map[string]lockfile.PackageDetails{}
+	collector := pyprojectPackageCollector{
+		packages:       map[string]lockfile.PackageDetails{},
+		lines:          lines,
+		path:           f.Path(),
+		packageManager: pm,
+	}
 
 	for _, dep := range pyproject.Project.Dependencies {
-		if name, rawName, version, ok := parsePEP508Pin(dep); ok {
-			block, nameLocation, versionLocation := extractPositions(lines, f.Path(), rawName, version, false)
-			addOrMergeGroups(packages, name, version, []string{"prod"}, pm, block, nameLocation, versionLocation)
+		dependency, ok := parsePEP508Dependency(dep)
+		if ok {
+			collector.addDependency(dependency, []string{"prod"}, false)
 		}
 	}
 
 	for group, deps := range pyproject.Project.OptionalDependencies {
 		for _, dep := range deps {
-			if name, rawName, version, ok := parsePEP508Pin(dep); ok {
-				block, nameLocation, versionLocation := extractPositions(lines, f.Path(), rawName, version, false)
-				addOrMergeGroups(packages, name, version, []string{group}, pm, block, nameLocation, versionLocation)
+			dependency, ok := parsePEP508Dependency(dep)
+			if ok {
+				collector.addDependency(dependency, []string{group}, false)
 			}
 		}
 	}
@@ -172,9 +178,9 @@ func (e PyProjectTOMLExtractor) Extract(f lockfile.DepFile, context lockfile.Sca
 				// skip {include-group = "..."} table entries
 				continue
 			}
-			if name, rawName, version, ok := parsePEP508Pin(dep); ok {
-				block, nameLocation, versionLocation := extractPositions(lines, f.Path(), rawName, version, false)
-				addOrMergeGroups(packages, name, version, []string{group}, pm, block, nameLocation, versionLocation)
+			dependency, ok := parsePEP508Dependency(dep)
+			if ok {
+				collector.addDependency(dependency, []string{group}, false)
 			}
 		}
 	}
@@ -185,147 +191,225 @@ func (e PyProjectTOMLExtractor) Extract(f lockfile.DepFile, context lockfile.Sca
 			if name == "python" {
 				continue
 			}
-			if version, ok := parsePoetryPin(val); ok {
-				normalized := normalizedRequirementName(name)
-				block, nameLocation, versionLocation := extractPositions(lines, f.Path(), name, version, true)
-				addOrMergeGroups(packages, normalized, version, []string{"prod"}, pm, block, nameLocation, versionLocation)
+			if version, versionRange, ok := parsePoetryDependency(val); ok {
+				collector.addDependency(pep508Dependency{
+					Name:         normalizedRequirementName(name),
+					RawName:      name,
+					Version:      version,
+					VersionRange: versionRange,
+				}, []string{"prod"}, true)
 			}
 		}
 		for name, val := range pyproject.Tool.Poetry.DevDependencies {
-			if version, ok := parsePoetryPin(val); ok {
-				normalized := normalizedRequirementName(name)
-				block, nameLocation, versionLocation := extractPositions(lines, f.Path(), name, version, true)
-				addOrMergeGroups(packages, normalized, version, []string{"dev"}, pm, block, nameLocation, versionLocation)
+			if version, versionRange, ok := parsePoetryDependency(val); ok {
+				collector.addDependency(pep508Dependency{
+					Name:         normalizedRequirementName(name),
+					RawName:      name,
+					Version:      version,
+					VersionRange: versionRange,
+				}, []string{"dev"}, true)
 			}
 		}
 		for groupName, group := range pyproject.Tool.Poetry.Group {
 			for name, val := range group.Dependencies {
-				if version, ok := parsePoetryPin(val); ok {
-					normalized := normalizedRequirementName(name)
-					block, nameLocation, versionLocation := extractPositions(lines, f.Path(), name, version, true)
-					addOrMergeGroups(packages, normalized, version, []string{groupName}, pm, block, nameLocation, versionLocation)
+				if version, versionRange, ok := parsePoetryDependency(val); ok {
+					collector.addDependency(pep508Dependency{
+						Name:         normalizedRequirementName(name),
+						RawName:      name,
+						Version:      version,
+						VersionRange: versionRange,
+					}, []string{groupName}, true)
 				}
 			}
 		}
 	}
 
-	return slices.Collect(maps.Values(packages)), nil
+	return slices.Collect(maps.Values(collector.packages)), nil
 }
 
-// addOrMergeGroups adds a package to the map, or if it already exists (same name+version),
-// merges the new dep groups into the existing entry rather than dropping the duplicate.
-func addOrMergeGroups(packages map[string]lockfile.PackageDetails, name, version string, groups []string, pm models.PackageManager, block models.FilePosition, nameLocation, versionLocation *models.FilePosition) {
-	key := name + "@" + version
-	if existing, exists := packages[key]; exists {
-		for _, g := range groups {
-			if !slices.Contains(existing.DepGroups, g) {
-				existing.DepGroups = append(existing.DepGroups, g)
-			}
-		}
-		packages[key] = existing
+type pyprojectPackageCollector struct {
+	packages       map[string]lockfile.PackageDetails
+	lines          []string
+	path           string
+	packageManager models.PackageManager
+}
+
+func (c *pyprojectPackageCollector) addDependency(dependency pep508Dependency, groups []string, isPoetry bool) {
+	if !dependency.hasExactlyOneVersionValue() {
+		log.Printf(
+			"Skipping pyproject dependency %q from %s: expected exactly one of version or version range, got version=%q versionRange=%q\n",
+			dependency.Name,
+			c.path,
+			dependency.Version,
+			dependency.VersionRange,
+		)
 
 		return
 	}
-	packages[key] = lockfile.PackageDetails{
-		Name:            name,
-		Version:         version,
-		PackageManager:  pm,
+
+	block, nameLocation, versionLocation := extractPositions(c.lines, c.path, dependency.RawName, versionOrRange(dependency.Version, dependency.VersionRange), isPoetry)
+	c.addOrMergePackageGroups(lockfile.PackageDetails{
+		Name:            dependency.Name,
+		Version:         dependency.Version,
+		VersionRange:    dependency.VersionRange,
+		PackageManager:  c.packageManager,
 		Ecosystem:       models.EcosystemPyPI,
 		IsDirect:        true,
 		DepGroups:       groups,
 		BlockLocation:   block,
 		NameLocation:    nameLocation,
 		VersionLocation: versionLocation,
 		LocationRole:    models.LocationRoleManifest,
+	})
+}
+
+// addOrMergePackageGroups adds a package to the map, or if it already exists (same name+version/range),
+// merges the new dep groups into the existing entry rather than dropping the duplicate.
+func (c *pyprojectPackageCollector) addOrMergePackageGroups(pkg lockfile.PackageDetails) {
+	key := pkg.Name + "@" + pkg.Version + "|" + pkg.VersionRange
+	if existing, exists := c.packages[key]; exists {
+		for _, g := range pkg.DepGroups {
+			if !slices.Contains(existing.DepGroups, g) {
+				existing.DepGroups = append(existing.DepGroups, g)
+			}
+		}
+		c.packages[key] = existing
+
+		return
 	}
+	c.packages[key] = pkg
 }
 
-// parsePEP508Pin parses a PEP 508 dependency string and returns the normalized name, the raw
-// (pre-normalization) name as written in the file, and the version — only when the dependency
-// is an exact pin (==). Returns ok=false for all other specifiers.
-func parsePEP508Pin(dep string) (name, rawName, version string, ok bool) {
+func versionOrRange(version, versionRange string) string {
+	if version != "" {
+		return version
+	}
+
+	return versionRange
+}
+
+type pep508Dependency struct {
+	Name         string
+	RawName      string
+	Version      string
+	VersionRange string
+}
+
+func (d pep508Dependency) hasExactlyOneVersionValue() bool {
+	return (d.Version == "") != (d.VersionRange == "")
+}
+
+// parsePEP508Dependency parses a PEP 508 dependency string into a normalized name,
+// the raw name as written in the file, and either an exact version or original version range.
+func parsePEP508Dependency(dep string) (pep508Dependency, bool) {
 	// strip environment markers (PEP 508)
 	dep, _, _ = strings.Cut(dep, ";")
 	dep = strings.TrimSpace(dep)
+	if strings.Contains(dep, " @ ") {
+		return pep508Dependency{}, false
+	}
 	// strip parenthesized specifier: "requests (==2.28.0)" -> "requests ==2.28.0"
 	dep = strings.NewReplacer("(", "", ")", "").Replace(dep)
 
-	// reject if any non-exact operator is present
-	for _, op := range []string{"===", "!=", ">=", "<=", "~=", ">", "<"} {
-		if strings.Contains(dep, op) {
-			return "", "", "", false
-		}
+	opIndex, op := findFirstPEP508Specifier(dep)
+	if opIndex == -1 || op == "===" {
+		return pep508Dependency{}, false
 	}
 
-	rawNamePart, rawVersion, found := strings.Cut(dep, "==")
-	if !found {
-		return "", "", "", false
+	// strip extras: "requests[security]" -> "requests"
+	fileRawName, _, _ := strings.Cut(strings.TrimSpace(dep[:opIndex]), "[")
+	fileRawName = strings.TrimSpace(fileRawName)
+	specifier := strings.TrimSpace(dep[opIndex:])
+
+	if fileRawName == "" || specifier == "" {
+		return pep508Dependency{}, false
 	}
 
-	// reject multi-constraint specs: "==2.28.0,!=2.28.0"
-	if strings.Contains(rawVersion, ",") {
-		return "", "", "", false
+	if op == "==" {
+		rawVersion := strings.TrimSpace(specifier[len(op):])
+		if rawVersion != "" && !strings.Contains(rawVersion, ",") && isConcreteVersion(rawVersion) {
+			return pep508Dependency{
+				Name:    normalizedRequirementName(fileRawName),
+				RawName: fileRawName,
+				Version: rawVersion,
+			}, true
+		}
 	}
-	rawVersion = strings.TrimSpace(rawVersion)
 
-	// strip extras: "requests[security]" -> "requests"
-	fileRawName, _, _ := strings.Cut(strings.TrimSpace(rawNamePart), "[")
-	fileRawName = strings.TrimSpace(fileRawName)
+	return pep508Dependency{
+		Name:         normalizedRequirementName(fileRawName),
+		RawName:      fileRawName,
+		VersionRange: specifier,
+	}, true
+}
 
-	if fileRawName == "" || rawVersion == "" || !isConcreteVersion(rawVersion) {
-		return "", "", "", false
+func findFirstPEP508Specifier(dep string) (int, string) {
+	firstIndex := -1
+	firstOp := ""
+	for _, op := range []string{"===", "==", "!=", ">=", "<=", "~=", ">", "<"} {
+		index := strings.Index(dep, op)
+		if index == -1 {
+			continue
+		}
+		if firstIndex == -1 || index < firstIndex {
+			firstIndex = index
+			firstOp = op
+		}
 	}
 
-	return normalizedRequirementName(fileRawName), fileRawName, rawVersion, true
+	return firstIndex, firstOp
 }
 
-// parsePoetryPin parses a Poetry dependency value (string or inline table) and returns
-// the version only when it is an exact pin (== prefix) with a concrete version.
-func parsePoetryPin(val any) (version string, ok bool) {
+// parsePoetryDependency parses a Poetry dependency value (string or inline table) and
+// returns either an exact version or the original version range.
+func parsePoetryDependency(val any) (version, versionRange string, ok bool) {
 	var versionStr string
 	switch v := val.(type) {
 	case string:
 		versionStr = v
 	case map[string]any:
+		for _, directRefKey := range []string{"path", "git", "url"} {
+			if _, exists := v[directRefKey]; exists {
+				return "", "", false
+			}
+		}
 		versionStr, ok = v["version"].(string)
 		if !ok {
-			return "", false
+			return "", "", false
 		}
 	default:
-		return "", false
+		return "", "", false
 	}
 
 	versionStr = strings.TrimSpace(versionStr)
+	if versionStr == "" || strings.HasPrefix(versionStr, "===") {
+		return "", "", false
+	}
 
-	// Poetry bare version string "2.28.0" is an implicit exact pin
+	// Poetry bare version string "2.28.0" is an implicit exact pin.
+	// Other digit-starting constraints, such as "1.*", are still ranges.
 	if len(versionStr) > 0 && !strings.ContainsAny(string(versionStr[0]), "=!<>~^*") {
 		if strings.Contains(versionStr, ",") {
-			return "", false
+			return "", versionStr, true
 		}
 		if isConcreteVersion(versionStr) {
-			return versionStr, true
+			return versionStr, "", true
 		}
 
-		return "", false
-	}
-
-	// reject === (arbitrary equality) and any non-== operator
-	if !strings.HasPrefix(versionStr, "==") || strings.HasPrefix(versionStr, "===") {
-		return "", false
+		return "", versionStr, true
 	}
 
-	after := strings.TrimSpace(versionStr[2:])
+	if strings.HasPrefix(versionStr, "==") {
+		after := strings.TrimSpace(versionStr[2:])
 
-	// reject multi-constraint: "==2.28.0,!=2.28.1" is not an exact pin
-	if strings.Contains(after, ",") {
-		return "", false
-	}
+		if !strings.Contains(after, ",") && isConcreteVersion(after) {
+			return after, "", true
+		}
 
-	if isConcreteVersion(after) {
-		return after, true
+		return "", versionStr, true
 	}
 
-	return "", false
+	return "", versionStr, true
 }
 
 // isConcreteVersion returns true if version looks like a fully-specified version
diff --git a/pkg/lockfile/python/parse-pyproject-toml_test.go b/pkg/lockfile/python/parse-pyproject-toml_test.go
