diff --git a/python/lib/dependabot/python/package/package_registry_finder.rb b/python/lib/dependabot/python/package/package_registry_finder.rb index 396bae16c0e..7c9873e0b68 100644 --- a/python/lib/dependabot/python/package/package_registry_finder.rb +++ b/python/lib/dependabot/python/package/package_registry_finder.rb @@ -59,6 +59,9 @@ def registry_urls sig { returns(T::Array[Dependabot::Credential]) } attr_reader :credentials + sig { returns(Dependabot::Dependency) } + attr_reader :dependency + sig { returns(String) } def main_index_url url = @@ -256,7 +259,17 @@ def pyproject sig { returns(T::Array[Dependabot::DependencyFile]) } def requirements_files - dependency_files.select { |f| f.name.match?(/requirements/x) } + # Get the list of requirement file names where this dependency appears + requirement_file_names = dependency.requirements + .map { |r| r[:file] } + .select { |f| f.match?(/requirements/x) } + + # If the dependency has no requirements in any requirements files, + # return all requirements files (backward compatibility) + return dependency_files.select { |f| f.name.match?(/requirements/x) } if requirement_file_names.empty? + + # Otherwise, only return the requirement files where this dependency appears + dependency_files.select { |f| requirement_file_names.include?(f.name) } end sig { returns(T::Array[Dependabot::DependencyFile]) } diff --git a/python/spec/dependabot/python/package/package_registry_finder_spec.rb b/python/spec/dependabot/python/package/package_registry_finder_spec.rb index b450d09fb4f..f29abc3c7ee 100644 --- a/python/spec/dependabot/python/package/package_registry_finder_spec.rb +++ b/python/spec/dependabot/python/package/package_registry_finder_spec.rb @@ -345,6 +345,102 @@ expect(registry_urls).to contain_exactly("https://pypi.org/simple/", "https://pypi.weasyldev.com/weasyl/source/+simple/") end end + + context "with multiple requirements files each having different extra-index-url" do + let(:requirements_cpu) do + Dependabot::DependencyFile.new( + name: "requirements_cpu.txt", + content: fixture("requirements", "torch_cpu.txt") + ) + end + let(:requirements_gpu) do + Dependabot::DependencyFile.new( + name: "requirements_gpu.txt", + content: fixture("requirements", "torch_gpu.txt") + ) + end + + context "when dependency is only in requirements_cpu.txt" do + let(:dependency_files) { [requirements_cpu, requirements_gpu] } + let(:dependency) do + Dependabot::Dependency.new( + name: "torch", + version: "2.0.0+cpu", + requirements: [{ + requirement: "==2.0.0+cpu", + file: "requirements_cpu.txt", + groups: ["dependencies"], + source: nil + }], + package_manager: "pip" + ) + end + + it "only uses the extra-index-url from requirements_cpu.txt" do + expect(registry_urls).to contain_exactly( + "https://pypi.org/simple/", + "https://download.pytorch.org/whl/cpu/" + ) + end + end + + context "when dependency is only in requirements_gpu.txt" do + let(:dependency_files) { [requirements_cpu, requirements_gpu] } + let(:dependency) do + Dependabot::Dependency.new( + name: "torch", + version: "2.0.0+cu118", + requirements: [{ + requirement: "==2.0.0+cu118", + file: "requirements_gpu.txt", + groups: ["dependencies"], + source: nil + }], + package_manager: "pip" + ) + end + + it "only uses the extra-index-url from requirements_gpu.txt" do + expect(registry_urls).to contain_exactly( + "https://pypi.org/simple/", + "https://download.pytorch.org/whl/cu118/" + ) + end + end + + context "when dependency is in both files" do + let(:dependency_files) { [requirements_cpu, requirements_gpu] } + let(:dependency) do + Dependabot::Dependency.new( + name: "torch", + version: "2.0.0+cpu", + requirements: [ + { + requirement: "==2.0.0+cpu", + file: "requirements_cpu.txt", + groups: ["dependencies"], + source: nil + }, + { + requirement: "==2.0.0+cu118", + file: "requirements_gpu.txt", + groups: ["dependencies"], + source: nil + } + ], + package_manager: "pip" + ) + end + + it "uses extra-index-url from both files" do + expect(registry_urls).to contain_exactly( + "https://pypi.org/simple/", + "https://download.pytorch.org/whl/cpu/", + "https://download.pytorch.org/whl/cu118/" + ) + end + end + end end end end diff --git a/python/spec/fixtures/requirements/torch_cpu.txt b/python/spec/fixtures/requirements/torch_cpu.txt new file mode 100644 index 00000000000..30e3683a920 --- /dev/null +++ b/python/spec/fixtures/requirements/torch_cpu.txt @@ -0,0 +1,2 @@ +--extra-index-url https://download.pytorch.org/whl/cpu +torch==2.0.0+cpu diff --git a/python/spec/fixtures/requirements/torch_gpu.txt b/python/spec/fixtures/requirements/torch_gpu.txt new file mode 100644 index 00000000000..8ec296fb4bd --- /dev/null +++ b/python/spec/fixtures/requirements/torch_gpu.txt @@ -0,0 +1,2 @@ +--extra-index-url https://download.pytorch.org/whl/cu118 +torch==2.0.0+cu118