wrap urlopen with requests

Author: ocefpaf

pandas/io/common.py

@@ -31,14 +31,13 @@
 
 if compat.PY3:
     from urllib.request import urlopen, pathname2url
-    _urlopen = urlopen
     from urllib.parse import urlparse as parse_url
     from urllib.parse import (uses_relative, uses_netloc, uses_params,
                               urlencode, urljoin)
     from urllib.error import URLError
     from http.client import HTTPException  # noqa
 else:
-    from urllib2 import urlopen as _urlopen
+    from urllib2 import urlopen as urlopen2
     from urllib import urlencode, pathname2url  # noqa
     from urlparse import urlparse as parse_url
     from urlparse import uses_relative, uses_netloc, uses_params, urljoin
@@ -47,10 +46,10 @@
     from contextlib import contextmanager, closing  # noqa
     from functools import wraps  # noqa
 
-    # @wraps(_urlopen)
+    # @wraps(urlopen2)
     @contextmanager
     def urlopen(*args, **kwargs):
-        with closing(_urlopen(*args, **kwargs)) as f:
+        with closing(urlopen2(*args, **kwargs)) as f:
             yield f
 
 
@@ -92,6 +91,34 @@ def _is_url(url):
         return False
 
 
+def _urlopen(url, session=None):
+    compression = None
+    content_encoding = None
+    try:
+        import requests
+        if session:
+            if not isinstance(session, requests.sessions.Session):
+                raise ValueError(
+                    'Expected a requests.sessions.Session object, '
+                    'got {!r}'.format(session)
+                )
+            r = session.get(url)
+        else:
+            r = requests.get(url)
+        r.raise_for_status
+        content = r.content
+    except ImportError:
+        r = urlopen(url)
+        content = r.read()
+        content_encoding = r.headers.get('Content-Encoding', None)
+    r.close()
+    if content_encoding == 'gzip':
+        # Override compression based on Content-Encoding header.
+        compression = 'gzip'
+    reader = BytesIO(content)
+    return reader, compression
+
+
 def _expand_user(filepath_or_buffer):
     """Return the argument with an initial component of ~ or ~user
        replaced by that user's home directory.
@@ -170,7 +197,7 @@ def is_s3_url(url):
 
 
 def get_filepath_or_buffer(filepath_or_buffer, encoding=None,
-                           compression=None, mode=None):
+                           compression=None, mode=None, session=None):
     """
     If the filepath_or_buffer is a url, translate and return the buffer.
     Otherwise passthrough.
@@ -192,13 +219,7 @@ def get_filepath_or_buffer(filepath_or_buffer, encoding=None,
     filepath_or_buffer = _stringify_path(filepath_or_buffer)
 
     if _is_url(filepath_or_buffer):
-        req = _urlopen(filepath_or_buffer)
-        content_encoding = req.headers.get('Content-Encoding', None)
-        if content_encoding == 'gzip':
-            # Override compression based on Content-Encoding header
-            compression = 'gzip'
-        reader = BytesIO(req.read())
-        req.close()
+        reader, compression = _urlopen(filepath_or_buffer, session=session)
         return reader, encoding, compression, True
 
     if is_s3_url(filepath_or_buffer):

pandas/io/parsers.py

@@ -297,6 +297,9 @@
     If a filepath is provided for `filepath_or_buffer`, map the file object
     directly onto memory and access the data directly from there. Using this
     option can improve performance because there is no longer any I/O overhead.
+session : requests.Session
+    object with the a requests session configuration for remote file.
+    (requires the requests library)
 
 Returns
 -------
@@ -418,10 +421,11 @@ def _read(filepath_or_buffer, kwds):
         encoding = re.sub('_', '-', encoding).lower()
         kwds['encoding'] = encoding
 
+    session = kwds.get('session', None)
     compression = kwds.get('compression')
     compression = _infer_compression(filepath_or_buffer, compression)
     filepath_or_buffer, _, compression, should_close = get_filepath_or_buffer(
-        filepath_or_buffer, encoding, compression)
+        filepath_or_buffer, encoding, compression, session=session)
     kwds['compression'] = compression
 
     if kwds.get('date_parser', None) is not None:
@@ -604,7 +608,8 @@ def parser_f(filepath_or_buffer,
                  delim_whitespace=False,
                  low_memory=_c_parser_defaults['low_memory'],
                  memory_map=False,
-                 float_precision=None):
+                 float_precision=None,
+                 session=None):
 
         # Alias sep -> delimiter.
         if delimiter is None:
@@ -673,7 +678,8 @@ def parser_f(filepath_or_buffer,
                     mangle_dupe_cols=mangle_dupe_cols,
                     tupleize_cols=tupleize_cols,
                     infer_datetime_format=infer_datetime_format,
-                    skip_blank_lines=skip_blank_lines)
+                    skip_blank_lines=skip_blank_lines,
+                    session=session)
 
         return _read(filepath_or_buffer, kwds)