Repository: streamlink/streamlink
Pull request #3746 (2021-05-23 19:14:13)

utils.url: add encoding options to update_qsd
Add `safe` and `quote_via` parameters to `update_qsd`, which are
passed to `urllib.parse.urlencode` for encoding the query string
keys and values.


Patch:
diff --git a/src/streamlink/utils/url.py b/src/streamlink/utils/url.py
index 35bbfb34..7776c648 100644
--- a/src/streamlink/utils/url.py
+++ b/src/streamlink/utils/url.py
@@ -1,5 +1,5 @@
 from collections import OrderedDict
-from urllib.parse import parse_qsl, urlencode, urljoin, urlparse, urlunparse
+from urllib.parse import parse_qsl, quote_plus, urlencode, urljoin, urlparse, urlunparse
 
 
 def update_scheme(current, target):
@@ -64,7 +64,7 @@ def url_concat(base, *parts, **kwargs):
     return base
 
 
-def update_qsd(url, qsd=None, remove=None, keep_blank_values=True):
+def update_qsd(url, qsd=None, remove=None, keep_blank_values=True, safe="", quote_via=quote_plus):
     """
     Update or remove keys from a query string in a URL
 
@@ -72,7 +72,9 @@ def update_qsd(url, qsd=None, remove=None, keep_blank_values=True):
     :param qsd: dict of keys to update, a None value leaves it unchanged
     :param remove: list of keys to remove, or "*" to remove all
                    note: updated keys are never removed, even if unchanged
-    :param keep_blank_values: if params with blank values should be kept or not
+    :param keep_blank_values: whether params with blank values should be kept or not
+    :param safe: string of reserved encoding characters, passed to the quote_via function
+    :param quote_via: function which encodes query string keys and values. Default: urllib.parse.quote_plus
     :return: updated URL
     """
     qsd = qsd or {}
@@ -100,4 +102,6 @@ def update_qsd(url, qsd=None, remove=None, keep_blank_values=True):
         if not value and not keep_blank_values and key not in qsd:
             del current_qsd[key]
 
-    return parsed._replace(query=urlencode(current_qsd)).geturl()
+    query = urlencode(query=current_qsd, safe=safe, quote_via=quote_via)
+
+    return parsed._replace(query=query).geturl()


Test patch:
diff --git a/tests/test_utils_url.py b/tests/test_utils_url.py
index eb87e6df..9f283316 100644
--- a/tests/test_utils_url.py
+++ b/tests/test_utils_url.py
@@ -1,4 +1,5 @@
 from collections import OrderedDict
+from urllib.parse import quote
 
 from streamlink.utils.url import update_qsd, update_scheme, url_concat, url_equal
 
@@ -43,3 +44,16 @@ def test_update_qsd():
     assert update_qsd("http://test.se?&two=", {"one": ''}, keep_blank_values=False) == "http://test.se?one=", \
         "should set one blank"
     assert update_qsd("http://test.se?one=", {"two": 2}) == "http://test.se?one=&two=2"
+
+    assert update_qsd("http://test.se?foo=%3F", {"bar": "!"}) == "http://test.se?foo=%3F&bar=%21", \
+        "urlencode - encoded URL"
+    assert update_qsd("http://test.se?foo=?", {"bar": "!"}) == "http://test.se?foo=%3F&bar=%21", \
+        "urlencode - fix URL"
+    assert update_qsd("http://test.se?foo=?", {"bar": "!"}, quote_via=lambda s, *_: s) == "http://test.se?foo=?&bar=!", \
+        "urlencode - dummy quote method"
+    assert update_qsd("http://test.se", {"foo": "/ "}) == "http://test.se?foo=%2F+", \
+        "urlencode - default quote_plus"
+    assert update_qsd("http://test.se", {"foo": "/ "}, safe="/", quote_via=quote) == "http://test.se?foo=/%20", \
+        "urlencode - regular quote with reserved slash"
+    assert update_qsd("http://test.se", {"foo": "/ "}, safe="", quote_via=quote) == "http://test.se?foo=%2F%20", \
+        "urlencode - regular quote without reserved slash"