refactor: improve download streaming and progress bar logic by extracting helper functions and enhancing readability.

Author: dankkom

comexdown/__init__.py

@@ -31,7 +31,12 @@ def get_year(path: Path, year: int, exp=False, imp=False, mun=False):
 
     for direction in directions:
         url = urls.trade(direction=direction, year=year, mun=mun)
-        file_path = fs.path_trade(root=path, direction=direction, year=year, mun=mun)
+        file_path = fs.path_trade(
+            root=path,
+            direction=direction,
+            year=year,
+            mun=mun,
+        )
         download.download_file(url, file_path)
 
 
@@ -57,7 +62,11 @@ def get_year_nbm(path: Path, year: int, exp=False, imp=False):
 
     for direction in directions:
         url = urls.trade(direction=direction, year=year, nbm=True)
-        file_path = fs.path_trade_nbm(root=path, direction=direction, year=year)
+        file_path = fs.path_trade_nbm(
+            root=path,
+            direction=direction,
+            year=year,
+        )
         download.download_file(url, file_path)
 
 
@@ -84,17 +93,22 @@ def get_complete(path: Path, exp=False, imp=False, mun=False):
     for direction in directions:
         url = urls.complete(direction=direction, mun=mun)
         # Note: 'complete' files might have different naming conventions
-        # The original code relied on download.exp_complete which hardcoded the filename.
-        # fs.path_trade generates paths like .../exp/EXP_2020.csv, which isn't right for complete zip files.
+        # The original code relied on download.exp_complete which hardcoded the
+        # filename.
+        # fs.path_trade generates paths like .../exp/EXP_2020.csv, which isn't
+        # right for complete zip files.
         # We need to handle the output path for complete files.
         # The original code did: path / filename (where filename is separate).
         # We need to replicate that logic or add it to fs.py.
-        # Let's simple determine the filename from the URL for now as the original did.
+        # Let's simple determine the filename from the URL for now as the
+        # original did.
         filename = url.split("/")[-1]
         file_path = path / filename
 
-        # Original code for complete files saved directly to `path` (or `path` was a directory).
-        # The original implementation for complete files: `filepath = path / filename`.
+        # Original code for complete files saved directly to `path` (or `path`
+        # was a directory).
+        # The original implementation for complete files:
+        # `filepath = path / filename`.
         # So we expect `path` to be a directory.
         download.download_file(url, file_path)
 

comexdown/download.py

@@ -10,7 +10,7 @@
 urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
 
 
-def is_more_recent(headers: dict, dest: Path) -> bool:
+def remote_is_more_recent(headers: dict, dest: Path) -> bool:
     """Check if the remote file is more recent than the local file."""
     if not dest.exists():
         return False
@@ -27,6 +27,42 @@ def is_more_recent(headers: dict, dest: Path) -> bool:
     return False
 
 
+def _print_progress(
+    downloaded: int,
+    total: int,
+    width: int = 50,
+) -> None:
+    if not total:
+        return
+
+    p = downloaded / total
+    filled = int(p * width)
+    bar = "=" * filled + "-" * (width - filled)
+    size_mb = downloaded / (1024 * 1024)
+    msg = f"\r[{bar}] {p:.1%} ({size_mb:.2f} MiB)"
+    sys.stdout.write(msg)
+    sys.stdout.flush()
+
+
+def _download_stream(
+    response: requests.Response,
+    output: Path,
+    blocksize: int,
+) -> None:
+    response.raise_for_status()
+    total_length = int(response.headers.get("content-length", 0))
+
+    downloaded_size = 0
+    with open(output, "wb") as f:
+        for chunk in response.iter_content(chunk_size=blocksize):
+            if chunk:
+                f.write(chunk)
+                downloaded_size += len(chunk)
+                _print_progress(downloaded_size, total_length)
+
+    sys.stdout.write("\n")
+
+
 def download_file(
     url: str,
     output: Path,
@@ -48,7 +84,11 @@ def download_file(
         The path to the downloaded file.
     """
     headers = {
-        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/142.0.0.0 Safari/537.36",
+        "User-Agent": (
+            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
+            "AppleWebKit/537.36 (KHTML, like Gecko) "
+            "Chrome/142.0.0.0 Safari/537.36"
+        ),
     }
 
     # Ensure parent directory exists
@@ -64,38 +104,23 @@ def download_file(
                 url, headers=headers, timeout=10, verify=verify_ssl
             )
 
-            if output.exists() and not is_more_recent(head_resp.headers, output):
-                sys.stdout.write(f"             {output.name} is up to date.\n")
+            # Check if local file is up to date (i.e. remote is NOT newer)
+            cond = remote_is_more_recent(head_resp.headers, output)
+            if output.exists() and not cond:
+                sys.stdout.write(f"{output.name} is up to date.\n")
                 sys.stdout.flush()
                 return output
 
             # Perform the specific download
             with requests.get(
-                url, headers=headers, stream=True, timeout=30, verify=verify_ssl
+                url,
+                headers=headers,
+                stream=True,
+                timeout=30,
+                verify=verify_ssl,
             ) as r:
-                r.raise_for_status()
-                total_length = int(r.headers.get("content-length", 0))
-
-                downloaded_size = 0
-                with open(output, "wb") as f:
-                    for chunk in r.iter_content(chunk_size=blocksize):
-                        if chunk:
-                            f.write(chunk)
-                            downloaded_size += len(chunk)
-
-                            # Simple progress bar
-                            if total_length:
-                                percent = downloaded_size / total_length
-                                bar_length = 50
-                                filled = int(percent * bar_length)
-                                bar = "=" * filled + "-" * (bar_length - filled)
-                                size_mb = downloaded_size / (1024 * 1024)
-                                sys.stdout.write(
-                                    f"\r[{bar}] {percent:.1%} ({size_mb:.2f} MiB)"
-                                )
-                                sys.stdout.flush()
-
-            sys.stdout.write("\n")
+                _download_stream(r, output, blocksize)
+
             return output
 
         except requests.RequestException as e:

comexdown/urls.py

@@ -9,7 +9,12 @@ def table(table_name: str) -> str:
     return f"{BASE_URL}tabelas/{tables.AUX_TABLES[table_name]}"
 
 
-def trade(direction: str, year: int, mun: bool = False, nbm: bool = False) -> str:
+def trade(
+    direction: str,
+    year: int,
+    mun: bool = False,
+    nbm: bool = False,
+) -> str:
     """
     Generates URL for trade data.
     direction: 'exp' or 'imp'

pyproject.toml

@@ -25,7 +25,6 @@ dependencies = [
     "requests>=2.32.5",
 ]
 optional-dependencies = { dev = [
-    "coveralls",
     "pytest",
     "pytest-cov",
     "pylint",

tests/test_download.py

@@ -22,7 +22,8 @@ def test_download_file(self, mock_open, mock_requests, mock_sys):
         mock_head.headers = {}
         mock_requests.head.return_value = mock_head
 
-        download.download_file("http://www.example.com/file.csv", Path("data/file.csv"))
+        download.download_file(
+            "http://www.example.com/file.csv", Path("data/file.csv"))
 
         mock_requests.get.assert_called()
         mock_open.assert_called_with(Path("data/file.csv"), "wb")