From 7db1c519192316b8445662cdf799f5b576aa505c Mon Sep 17 00:00:00 2001
From: "Paul Gauthier (aider)" <aider@paulg.org>
Date: Tue, 9 Jul 2024 16:00:29 +0100
Subject: [PATCH] Subprocess output was processed line-by-line to find and
 extract the PyTorch URL, and the subprocess was terminated as soon as the URL
 was found.

---
 setup.py | 23 +++++++++++++++--------
 1 file changed, 15 insertions(+), 8 deletions(-)

diff --git a/setup.py b/setup.py
index 7c97a941d..ecdb7f573 100644
--- a/setup.py
+++ b/setup.py
@@ -29,19 +29,26 @@ cmd = [
     "--no-deps",
     "--dest",
     "/dev/null",
-    # "--no-cache-dir",
-    # "--ignore-installed",
     "--index-url",
     "https://download.pytorch.org/whl/cpu",
 ]
 
-result = subprocess.check_output(cmd, text=True)
-print(result)
+pytorch_url = None
+try:
+    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
+    for line in process.stdout:
+        print(line, end='')  # Print each line of output
+        url_match = re.search(r"Downloading (https://download\.pytorch\.org/[^\s]+\.whl)", line)
+        if url_match:
+            pytorch_url = url_match.group(1)
+            print(f"PyTorch URL: {pytorch_url}")
+            process.terminate()  # Terminate the subprocess
+            break
+    process.wait()  # Wait for the process to finish
+except subprocess.CalledProcessError as e:
+    print(f"Error running pip download: {e}")
 
-url_match = re.search(r"Downloading (https://download\.pytorch\.org/[^\s]+\.whl)", result)
-if url_match:
-    pytorch_url = url_match.group(1)
-    print(f"PyTorch URL: {pytorch_url}")
+if pytorch_url:
     requirements = [f"torch @ {pytorch_url}"] + requirements
 else:
     print("PyTorch URL not found in the output")