From 87001a2778b563bb20bda5254b19c4da6e08c063 Mon Sep 17 00:00:00 2001
From: Ryan Miguel <rmiguel@barracuda.com>
Date: Tue, 18 Oct 2022 16:03:57 -0700
Subject: [PATCH] Adds a special case for index urls similar to how comments
 are captured. This should avoid breaking ordering when --index-url (-i)
 and/or --extra-index-url are used. Also adds a relevant test to cover this
 case. Closes #612.

---
 pre_commit_hooks/requirements_txt_fixer.py | 12 ++++++++++++
 tests/requirements_txt_fixer_test.py       | 21 +++++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/pre_commit_hooks/requirements_txt_fixer.py b/pre_commit_hooks/requirements_txt_fixer.py
index 5884394..0e4f75a 100644
--- a/pre_commit_hooks/requirements_txt_fixer.py
+++ b/pre_commit_hooks/requirements_txt_fixer.py
@@ -17,6 +17,7 @@ class Requirement:
     def __init__(self) -> None:
         self.value: bytes | None = None
         self.comments: list[bytes] = []
+        self.index_urls: list[bytes] = []
 
     @property
     def name(self) -> bytes:
@@ -64,6 +65,8 @@ def fix_requirements(f: IO[bytes]) -> int:
     requirements: list[Requirement] = []
     before = list(f)
     after: list[bytes] = []
+    index_options: list[bytes] = []
+    extra_index_options: list[bytes] = []
 
     before_string = b''.join(before)
 
@@ -97,6 +100,13 @@ def fix_requirements(f: IO[bytes]) -> int:
                 requirement.comments.append(line)
         elif line.lstrip().startswith(b'#') or line.strip() == b'':
             requirement.comments.append(line)
+        elif (
+                line.lstrip().startswith(b'-i ') or
+                line.lstrip().startswith(b'--index-url')
+        ):
+            index_options.append(line)
+        elif line.lstrip().startswith(b'--extra-index-url'):
+            extra_index_options.append(line)
         else:
             requirement.append_value(line)
 
@@ -113,6 +123,8 @@ def fix_requirements(f: IO[bytes]) -> int:
         if req.value != b'pkg-resources==0.0.0\n'
     ]
 
+    after.extend(index_options)
+    after.extend(extra_index_options)
     for requirement in sorted(requirements):
         after.extend(requirement.comments)
         assert requirement.value, requirement.value
diff --git a/tests/requirements_txt_fixer_test.py b/tests/requirements_txt_fixer_test.py
index b725afa..f9d941d 100644
--- a/tests/requirements_txt_fixer_test.py
+++ b/tests/requirements_txt_fixer_test.py
@@ -99,6 +99,27 @@ from pre_commit_hooks.requirements_txt_fixer import Requirement
             PASS,
             b'a=2.0.0 \\\n --hash=sha256:abcd\nb==1.0.0\n',
         ),
+        (
+            b'# Bar\n'
+            b'--extra-index-url http://dist.repoze.org/zope2/2.10/simple\n'
+            b'zopelib2\n'
+            b'# Foo\n'
+            b'-i http://dist.repoze.org/zope2/2.10/simple\n'
+            b'zopelib1\n'
+            b'# Baz\n'
+            b'--index-url http://dist.repoze.org/zope2/2.10/simple\n'
+            b'zopelib3',
+            FAIL,
+            b'-i http://dist.repoze.org/zope2/2.10/simple\n'
+            b'--index-url http://dist.repoze.org/zope2/2.10/simple\n'
+            b'--extra-index-url http://dist.repoze.org/zope2/2.10/simple\n'
+            b'# Foo\n'
+            b'zopelib1\n'
+            b'# Bar\n'
+            b'zopelib2\n'
+            b'# Baz\n'
+            b'zopelib3\n',
+        ),
     ),
 )
 def test_integration(input_s, expected_retval, output, tmpdir):