pyosmium tools: set timeouts via requests

Author: lonvia

src/osmium/replication/server.py

@@ -36,9 +36,14 @@ class ReplicationServer:
         internally keeps a connection to the server making downloads faster.
     """
 
-    def __init__(self, url: str, diff_type: str = 'osc.gz') -> None:
+    def __init__(self, url: str, diff_type: str = 'osc.gz',
+                 extra_request_params: Optional[Mapping[str, Any]] = None) -> None:
         self.baseurl = url
         self.diff_type = diff_type
+        if extra_request_params is None:
+            self.extra_request_params = dict(timeout=60, stream=True)
+        else:
+            self.extra_request_params = extra_request_params
         self.session: Optional[requests.Session] = None
 
     def close(self) -> None:
@@ -77,17 +82,19 @@ def open_url(self, url: urlrequest.Request) -> Any:
                 svr = ReplicationServer()
                 svr.open_url = opener.open
         """
-        headers = dict()
-        for h in url.header_items():
-            headers[h[0]] = h[1]
+        if 'headers' in self.extra_request_params:
+            get_params = self.extra_request_params
+        else:
+            get_params = dict(self.extra_request_params)
+            get_params['headers'] = {k: v for k,v in url.header_items()}
 
         if self.session is not None:
-            return self.session.get(url.get_full_url(), headers=headers, stream=True)
+            return self.session.get(url.get_full_url(), **get_params)
 
         @contextmanager
         def _get_url_with_session() -> Iterator[requests.Response]:
             with requests.Session() as session:
-                request = session.get(url.get_full_url(), headers=headers, stream=True)
+                request = session.get(url.get_full_url(), **get_params)
                 yield request
 
         return _get_url_with_session()

tools/pyosmium-get-changes

@@ -27,7 +27,6 @@ cookies to the server and will save received cookies to the jar file.
 
 from argparse import ArgumentParser, RawDescriptionHelpFormatter, ArgumentTypeError
 import datetime as dt
-import socket
 from osmium.replication import server as rserv
 from osmium.replication import newest_change_from_file
 from osmium.replication.utils import get_replication_header
@@ -211,9 +210,10 @@ def main(args):
             or 'https://planet.osm.org/replication/minute/'
     logging.info("Using replication server at %s" % url)
 
-    socket.setdefaulttimeout(options.socket_timeout)
+    extra_request_params = dict(stream=True,
+                                timeout=options.socket_timeout or None)
 
-    with rserv.ReplicationServer(url) as svr:
+    with rserv.ReplicationServer(url, extra_request_params=extra_request_params) as svr:
         if options.cookie is not None:
             # According to the documentation, the cookie jar loads the file only if FileCookieJar.load is called.
             cookie_jar = cookiejarlib.MozillaCookieJar(options.cookie)

tools/pyosmium-up-to-date

@@ -35,7 +35,6 @@ import re
 import sys
 import traceback
 import logging
-import socket
 
 from argparse import ArgumentParser, RawDescriptionHelpFormatter
 import datetime as dt
@@ -70,7 +69,10 @@ def update_from_osm_server(ts, options):
 
 def update_from_custom_server(url, seq, ts, options):
     """Update from a custom URL, simply using the diff sequence as is."""
-    with rserv.ReplicationServer(url, "osc.gz") as svr:
+    extra_request_params = dict(stream=True,
+                                timeout=options.socket_timeout or None)
+
+    with rserv.ReplicationServer(url, "osc.gz", extra_request_params=extra_request_params) as svr:
         if options.cookie is not None:
             # According to the documentation, the cookie jar loads the file only if FileCookieJar.load is called.
             cookie_jar = cookiejarlib.MozillaCookieJar(options.cookie)
@@ -244,8 +246,6 @@ if __name__ == '__main__':
     options = get_arg_parser(from_main=True).parse_args()
     log.setLevel(max(3 - options.loglevel, 0) * 10)
 
-    socket.setdefaulttimeout(options.socket_timeout)
-
     try:
         url, seq, ts = compute_start_point(options)
     except RuntimeError as e: