add request.update_query_params which update query params of given url and return new url

pydu/compat.py

@@ -17,10 +17,12 @@ if PY2:
     import urllib as urlib
     import urlparse
     from urlparse import urljoin
+    from urllib import urlencode
 else:
     import urllib.request as urlib
     import urllib.parse as urlparse
     from urllib.parse import urljoin
+    from urllib.parse import urlencode
 
 # Dictionary iteration
 if PY2:

pydu/request.py

@@ -5,7 +5,7 @@ import socket
 
 from . import logger
 from .string import safeunicode
-from .compat import PY2, string_types, urlparse, urlib
+from .compat import PY2, string_types, urlparse, urlib, urlencode
 
 
 class FileName(object):
@@ -125,3 +125,15 @@ def check_connect(ip, port, retry=1, timeout=0.5):
         finally:
             retry -= 1
     return None
+
+
+def update_query_params(url, params):
+    """
+    Update query params of given url and return new url.
+    """
+    parts = list(urlparse.urlparse(url))
+    query = dict(urlparse.parse_qsl(parts[4]))
+    query.update(params)
+    parts[4] = urlencode(query)
+    new_url = urlparse.urlunparse(parts)
+    return new_url

tests/test_request.py

@@ -2,7 +2,7 @@ import socket
 from .testing import mockserver
 import pydu.request
 from pydu.network import get_free_port
-from pydu.request import FileName, check_connect
+from pydu.request import FileName, check_connect, update_query_params
 
 
 def test_filename_from_url():
@@ -52,3 +52,9 @@ def test_check_connect(port=None):
 
     pydu.request.socket.socket = mock_socket
     assert not check_connect('127.0.0.1', port=port, timeout=0.01)
+
+
+def test_update_query_params():
+    assert update_query_params('http://example.com/', {'foo': 1}) == 'http://example.com/?foo=1'
+    assert update_query_params('http://example.com/?foo=1', {'foo': 2}) == 'http://example.com/?foo=2'
+    assert update_query_params('http://example.com/?foo=1', {'foo': 2, 'bar': 3}) == 'http://example.com/?foo=2&bar=3'