FIX: do not write failed downloads to disk in fetch_nb_dependencies

quantecon/util/notebooks.py

@@ -11,32 +11,30 @@
 
 "https://github.com/QuantEcon/QuantEcon.notebooks/raw/master/dependencies/mpi/something.py" --> ./something.py
 
-TODO
-----
-1. Write Style guide for QuantEcon.notebook contributions
-2. Write an interface for Dat Server
-3. Platform Agnostic (replace wget usage)
-
 """
 
 import os
+import warnings
 
 #-Remote Structure-#
 REPO = "https://github.com/QuantEcon/QuantEcon.notebooks"
 RAW = "raw"
 BRANCH = "master"
 #Hard Coded Dependencies Folder on QuantEcon.notebooks
 FOLDER = "dependencies"
+#-Default timeout (seconds) for the remote request: (connect, read)-#
+TIMEOUT = (10, 30)
 
 
-def fetch_nb_dependencies(files, repo=REPO, raw=RAW, branch=BRANCH, folder=FOLDER, overwrite=False, verbose=True):
+def fetch_nb_dependencies(files, repo=REPO, raw=RAW, branch=BRANCH, folder=FOLDER, overwrite=False,
+                          verbose=True, timeout=TIMEOUT):
     """
     Retrieve raw files from QuantEcon.notebooks or other Github repo
 
     Parameters
     ----------
-    file_list   list or dict
-                A list of files to specify a collection of filenames	
+    file_list   list, tuple or dict
+                A list or tuple of files to specify a collection of filenames
                 A dict of dir : list(files) to specify a directory
     repo        str, optional(default=REPO)
     raw 		str, optional(default=RAW)
@@ -45,6 +43,17 @@ def fetch_nb_dependencies(files, repo=REPO, raw=RAW, branch=BRANCH, folder=FOLDE
     folder      str, optional(default=FOLDER)
     overwrite   bool, optional(default=False)
     verbose     bool, optional(default=True)
+    timeout     float or tuple, optional(default=TIMEOUT)
+                Passed through to ``requests.get``. Either a single value or a
+                ``(connect, read)`` pair, in seconds.
+
+    Returns
+    -------
+    status      list(bool)
+                One entry per requested file: ``True`` if the file was fetched
+                and written, ``False`` if it was skipped or could not be
+                retrieved. A failed request emits a warning and leaves no file
+                on disk.
 
     Examples
     --------
@@ -74,8 +83,8 @@ def fetch_nb_dependencies(files, repo=REPO, raw=RAW, branch=BRANCH, folder=FOLDE
     import requests
 
     #-Generate Common Data Structure-#
-    if type(files) == list:
-        files = {"" : files}
+    if isinstance(files, (list, tuple)):
+        files = {"" : list(files)}
 
     status = []
 
@@ -98,9 +107,16 @@ def fetch_nb_dependencies(files, repo=REPO, raw=RAW, branch=BRANCH, folder=FOLDE
             if verbose: print("Fetching file: %s"%fl)
             #-Get file in OS agnostic way using requests-#
             url = "/".join([repo, raw, branch, folder, fl])
-            r = requests.get(url)
-            with open(fl, "wb") as fl:
-                fl.write(r.content)
+            try:
+                r = requests.get(url, timeout=timeout)
+                r.raise_for_status()
+            except requests.RequestException as e:
+                #-Do not write the response body: an error page is not the requested file-#
+                warnings.warn("Failed to fetch %s: %s" % (url, e), stacklevel=2)
+                status.append(False)
+                continue
+            with open(fl, "wb") as f:
+                f.write(r.content)
             status.append(True)
 
     return status

quantecon/util/tests/test_notebooks.py

@@ -8,7 +8,13 @@
 """
 
 from quantecon.util import fetch_nb_dependencies
+from quantecon.util.notebooks import TIMEOUT
+import http.server
 import os
+import socketserver
+import threading
+
+import pytest
 
 FILES = ['test_file.md']
 REPO = "https://github.com/QuantEcon/QuantEcon.py"
@@ -39,3 +45,82 @@ def test_fetch_nb_dependencies_overwrite(self):
 
     def teardown_method(self):
         os.remove("test_file.md")
+
+
+class _Handler(http.server.BaseHTTPRequestHandler):
+    """Serve one known file and 404 everything else."""
+
+    def do_GET(self):
+        if self.path.endswith("present.csv"):
+            body, code = b"a,b\n1,2\n", 200
+        else:
+            body, code = b"<html><title>404 Not Found</title></html>", 404
+        self.send_response(code)
+        self.send_header("Content-Length", str(len(body)))
+        self.end_headers()
+        self.wfile.write(body)
+
+    def log_message(self, *args):
+        pass
+
+
+@pytest.fixture
+def local_server():
+    server = socketserver.TCPServer(("127.0.0.1", 0), _Handler)
+    threading.Thread(target=server.serve_forever, daemon=True).start()
+    yield "http://127.0.0.1:{}".format(server.server_address[1])
+    server.shutdown()
+    server.server_close()
+
+
+class TestFetchFailures:
+    """A failed request must not be written to disk or reported as success."""
+
+    def test_missing_remote_file_is_not_written(self, local_server, tmp_path, monkeypatch):
+        monkeypatch.chdir(tmp_path)
+        with pytest.warns(UserWarning):
+            status = fetch_nb_dependencies(
+                files=['absent.csv'], repo=local_server, verbose=False)
+        assert status == [False]
+        assert not os.path.exists("absent.csv")
+
+    def test_successful_fetch_writes_content(self, local_server, tmp_path, monkeypatch):
+        monkeypatch.chdir(tmp_path)
+        status = fetch_nb_dependencies(
+            files=['present.csv'], repo=local_server, verbose=False)
+        assert status == [True]
+        with open("present.csv", "rb") as f:
+            assert f.read() == b"a,b\n1,2\n"
+
+    def test_tuple_input_is_accepted(self, local_server, tmp_path, monkeypatch):
+        """A tuple must behave like a list rather than falling through to the dict branch."""
+        monkeypatch.chdir(tmp_path)
+        status = fetch_nb_dependencies(
+            files=('present.csv',), repo=local_server, verbose=False)
+        assert status == [True]
+
+    def test_timeout_is_forwarded_to_requests(self, local_server, tmp_path, monkeypatch):
+        """The request must be bounded so a stalled server cannot block forever."""
+        import requests
+
+        seen = {}
+        original_get = requests.get
+
+        def spy(url, **kwargs):
+            seen.update(kwargs)
+            return original_get(url, **kwargs)
+
+        monkeypatch.chdir(tmp_path)
+        monkeypatch.setattr(requests, "get", spy)
+
+        #-An explicit timeout is forwarded-#
+        fetch_nb_dependencies(
+            files=['present.csv'], repo=local_server, verbose=False, timeout=5)
+        assert seen.get("timeout") == 5
+
+        #-And the default is a real bound, not None-#
+        seen.clear()
+        fetch_nb_dependencies(
+            files=['present.csv'], repo=local_server, verbose=False, overwrite=True)
+        assert seen.get("timeout") == TIMEOUT
+        assert TIMEOUT is not None