fix: implement resilient HTTP Range resume download logic in utilitie…

funannotate2/utilities.py

@@ -242,10 +242,62 @@ def download(url, name, wget=False, timeout=60, retries=3):
     # Try HTTPS first
     while attempt < retries:
         try:
+            headers = {}
+            existing_size = 0
+            if os.path.exists(file_name):
+                existing_size = os.path.getsize(file_name)
+
+            # Send HEAD request to check remote size and Range support
+            remote_size = None
+            try:
+                # Use a smaller timeout for HEAD to avoid hanging
+                with requests.head(url, timeout=10, allow_redirects=True, verify=False) as head_resp:
+                    if head_resp.status_code == 200:
+                        remote_size = int(head_resp.headers.get("Content-Length", 0))
+            except Exception:
+                pass
+
+            # If local file matches remote size, download is complete
+            if remote_size and existing_size == remote_size:
+                return
+
+            # If local file is larger than remote size, overwrite from scratch
+            if remote_size and existing_size > remote_size:
+                try:
+                    os.remove(file_name)
+                except Exception:
+                    pass
+                existing_size = 0
+
+            # If we have a partially downloaded file, try range request
+            if existing_size > 0:
+                headers["Range"] = f"bytes={existing_size}-"
+
             # Use requests for HTTP/HTTPS
-            with requests.get(url, stream=True, timeout=timeout, verify=False) as r:
+            with requests.get(url, stream=True, timeout=timeout, verify=False, headers=headers) as r:
+                if r.status_code == 206:
+                    write_mode = "ab"
+                elif r.status_code == 416:
+                    # Range Not Satisfiable (might be fully downloaded or server error)
+                    # Delete the file and start from scratch to be safe
+                    try:
+                        os.remove(file_name)
+                    except Exception:
+                        pass
+                    existing_size = 0
+                    headers.pop("Range", None)
+                    with requests.get(url, stream=True, timeout=timeout, verify=False) as r_retry:
+                        r_retry.raise_for_status()
+                        with open(file_name, "wb") as f:
+                            for chunk in r_retry.iter_content(chunk_size=8192):
+                                if chunk:
+                                    f.write(chunk)
+                    return
+                else:
+                    write_mode = "wb"
+
                 r.raise_for_status()  # Raise an exception for HTTP errors
-                with open(file_name, "wb") as f:
+                with open(file_name, write_mode) as f:
                     for chunk in r.iter_content(chunk_size=8192):
                         if chunk:  # filter out keep-alive chunks
                             f.write(chunk)
@@ -270,7 +322,10 @@ def download(url, name, wget=False, timeout=60, retries=3):
                         return
 
                 # If we get here, both HTTPS and FTP failed
-                raise Exception(f"Download failed: {str(e)}")
+                raise Exception(
+                    f"Download failed: {str(e)}. If you have unstable internet or download timeouts, "
+                    "try running with the '--wget' option to use wget for robust downloads."
+                )
 
 
 def _download_ftp(url, file_name, timeout=60):

tests/unit/test_utilities.py

@@ -8,6 +8,7 @@
 
 from funannotate2.utilities import (
     create_tmpdir,
+    download,
     merge_coordinates,
     naming_slug,
     readBlocks,
@@ -259,3 +260,87 @@ def test_monitoring_does_not_debug_log_formatted_memory_report(
         debug_messages = [call.args[0] for call in logger.debug.call_args_list]
         assert "snap input.fasta" in debug_messages
         assert not any("Memory usage for" in message for message in debug_messages)
+
+
+class TestDownload:
+    """Tests for the download function."""
+
+    @patch("funannotate2.utilities.requests.get")
+    @patch("funannotate2.utilities.requests.head")
+    def test_download_from_scratch(self, mock_head, mock_get, tmp_path):
+        """Test downloading a file from scratch."""
+        dest = tmp_path / "test_file.txt"
+
+        # Mock requests.head
+        mock_head_resp = MagicMock()
+        mock_head_resp.status_code = 200
+        mock_head_resp.headers = {"Content-Length": "10"}
+        mock_head.return_value.__enter__.return_value = mock_head_resp
+
+        # Mock requests.get
+        mock_get_resp = MagicMock()
+        mock_get_resp.status_code = 200
+        mock_get_resp.iter_content.return_value = [b"0123456789"]
+        mock_get.return_value.__enter__.return_value = mock_get_resp
+
+        download("https://example.com/file.txt", str(dest))
+
+        assert dest.exists()
+        assert dest.read_bytes() == b"0123456789"
+        mock_get.assert_called_once_with(
+            "https://example.com/file.txt",
+            stream=True,
+            timeout=60,
+            verify=False,
+            headers={},
+        )
+
+    @patch("funannotate2.utilities.requests.get")
+    @patch("funannotate2.utilities.requests.head")
+    def test_download_resume(self, mock_head, mock_get, tmp_path):
+        """Test resuming a partial download using range requests."""
+        dest = tmp_path / "test_file.txt"
+        dest.write_bytes(b"01234")  # Partial file of 5 bytes
+
+        # Mock requests.head
+        mock_head_resp = MagicMock()
+        mock_head_resp.status_code = 200
+        mock_head_resp.headers = {"Content-Length": "10"}
+        mock_head.return_value.__enter__.return_value = mock_head_resp
+
+        # Mock requests.get (returns range)
+        mock_get_resp = MagicMock()
+        mock_get_resp.status_code = 206
+        mock_get_resp.iter_content.return_value = [b"56789"]
+        mock_get.return_value.__enter__.return_value = mock_get_resp
+
+        download("https://example.com/file.txt", str(dest))
+
+        assert dest.exists()
+        assert dest.read_bytes() == b"0123456789"
+        mock_get.assert_called_once_with(
+            "https://example.com/file.txt",
+            stream=True,
+            timeout=60,
+            verify=False,
+            headers={"Range": "bytes=5-"},
+        )
+
+    @patch("funannotate2.utilities.requests.get")
+    @patch("funannotate2.utilities.requests.head")
+    def test_download_already_complete(self, mock_head, mock_get, tmp_path):
+        """Test that if the file is already complete, no download request is made."""
+        dest = tmp_path / "test_file.txt"
+        dest.write_bytes(b"0123456789")  # 10 bytes
+
+        # Mock requests.head
+        mock_head_resp = MagicMock()
+        mock_head_resp.status_code = 200
+        mock_head_resp.headers = {"Content-Length": "10"}
+        mock_head.return_value.__enter__.return_value = mock_head_resp
+
+        download("https://example.com/file.txt", str(dest))
+
+        # Check that get was never called
+        mock_get.assert_not_called()
+        assert dest.read_bytes() == b"0123456789"