diff --git a/quantecon/util/notebooks.py b/quantecon/util/notebooks.py index 5f73a654..8f96e3b7 100644 --- a/quantecon/util/notebooks.py +++ b/quantecon/util/notebooks.py @@ -11,15 +11,10 @@ "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" @@ -27,16 +22,19 @@ 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 diff --git a/quantecon/util/tests/test_notebooks.py b/quantecon/util/tests/test_notebooks.py index 56d0b195..b88c5c80 100644 --- a/quantecon/util/tests/test_notebooks.py +++ b/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"404 Not Found", 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