From 80e5740ca282ab938821c9adb71a9122542c68e9 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sun, 8 Sep 2024 15:47:36 -0500 Subject: [PATCH] Add extend_query method (#1128) --- CHANGES/1128.feature.rst | 3 ++ docs/api.rst | 42 ++++++++++++++++++++ tests/test_update_query.py | 78 ++++++++++++++++++++++++++++++++++++++ yarl/_url.py | 40 ++++++++++++++++++- 4 files changed, 162 insertions(+), 1 deletion(-) create mode 100644 CHANGES/1128.feature.rst diff --git a/CHANGES/1128.feature.rst b/CHANGES/1128.feature.rst new file mode 100644 index 000000000..2d7b4cd06 --- /dev/null +++ b/CHANGES/1128.feature.rst @@ -0,0 +1,3 @@ +Added :meth:`URL.extend_query() ` method, which can be used to extend parameters without replacing same named keys -- by :user:`bdraco`. + +This method was primarily added to replace the inefficient hand rolled method currently used in ``aiohttp``. diff --git a/docs/api.rst b/docs/api.rst index 05a4db613..b9e950b52 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -659,6 +659,48 @@ section generates a new :class:`URL` instance. Support subclasses of :class:`int` (except :class:`bool`) and :class:`float` as a query parameter value. +.. method:: URL.extend_query(query) + URL.extend_query(**kwargs) + + Returns a new URL with *query* part extended. + + Unlike :meth:`update_query`, this method keeps duplicate keys. + + Returned :class:`URL` object will contain query string which extends + parts from passed query parts (or parts of parsed query string). + + Accepts any :class:`~collections.abc.Mapping` (e.g. :class:`dict`, + :class:`~multidict.MultiDict` instances) or :class:`str`, + auto-encode the argument if needed. + + A sequence of ``(key, value)`` pairs is supported as well. + + Also it can take an arbitrary number of keyword arguments. + + Returns the same :class:`URL` if *query* of ``None`` is passed. + + .. note:: + + The library accepts :class:`str`, :class:`float`, :class:`int` and their + subclasses except :class:`bool` as query argument values. + + If a mapping such as :class:`dict` is used, the values may also be + :class:`list` or :class:`tuple` to represent a key has many values. + + Please see :ref:`yarl-bools-support` for the reason why :class:`bool` is not + supported out-of-the-box. + + .. doctest:: + + >>> URL('http://example.com/path?a=b&b=1').extend_query(b='2') + URL('http://example.com/path?a=b&b=1&b=2') + >>> URL('http://example.com/path?a=b&b=1').extend_query([('b', '2')]) + URL('http://example.com/path?a=b&b=1&b=2') + >>> URL('http://example.com/path?a=b&c=e&c=f').extend_query(c='d') + URL('http://example.com/path?a=b&c=e&c=f&c=d') + + .. versionadded:: 1.11.0 + .. method:: URL.update_query(query) URL.update_query(**kwargs) diff --git a/tests/test_update_query.py b/tests/test_update_query.py index 5756fbd67..26ccda62a 100644 --- a/tests/test_update_query.py +++ b/tests/test_update_query.py @@ -374,3 +374,81 @@ def test_update_query_with_mod_operator(): assert str(url % {"a": "1"} % {"b": "2"}) == "http://example.com/?a=1&b=2" assert str(url % {"a": "1"} % {"a": "3", "b": "2"}) == "http://example.com/?a=3&b=2" assert str(url / "foo" % {"a": "1"}) == "http://example.com/foo?a=1" + + +def test_extend_query(): + url = URL("http://example.com/") + assert str(url.extend_query({"a": "1"})) == "http://example.com/?a=1" + assert str(URL("test").extend_query(a=1)) == "test?a=1" + + url = URL("http://example.com/?foo=bar") + expected_url = URL("http://example.com/?foo=bar&baz=foo") + + assert url.extend_query({"baz": "foo"}) == expected_url + assert url.extend_query(baz="foo") == expected_url + assert url.extend_query("baz=foo") == expected_url + + +def test_extend_query_with_args_and_kwargs(): + url = URL("http://example.com/") + + with pytest.raises(ValueError): + url.extend_query("a", foo="bar") + + +def test_extend_query_with_multiple_args(): + url = URL("http://example.com/") + + with pytest.raises(ValueError): + url.extend_query("a", "b") + + +def test_extend_query_with_none_arg(): + url = URL("http://example.com/?foo=bar&baz=foo") + assert url.extend_query(None) == url + + +def test_extend_query_with_empty_dict(): + url = URL("http://example.com/?foo=bar&baz=foo") + assert url.extend_query({}) == url + + +def test_extend_query_existing_keys(): + url = URL("http://example.com/?a=2") + assert str(url.extend_query({"a": "1"})) == "http://example.com/?a=2&a=1" + assert str(URL("test").extend_query(a=1)) == "test?a=1" + + url = URL("http://example.com/?foo=bar&baz=original") + expected_url = URL("http://example.com/?foo=bar&baz=original&baz=foo") + + assert url.extend_query({"baz": "foo"}) == expected_url + assert url.extend_query(baz="foo") == expected_url + assert url.extend_query("baz=foo") == expected_url + + +def test_extend_query_with_args_and_kwargs_with_existing(): + url = URL("http://example.com/?a=original") + + with pytest.raises(ValueError): + url.extend_query("a", foo="bar") + + +def test_extend_query_with_non_ascii(): + url = URL("http://example.com/?foo=bar&baz=foo") + expected = URL("http://example.com/?foo=bar&baz=foo&%F0%9D%95%A6=%F0%9D%95%A6") + assert url.extend_query({"𝕦": "𝕦"}) == expected + + +def test_extend_query_with_non_ascii_as_str(): + url = URL("http://example.com/?foo=bar&baz=foo&") + expected = URL("http://example.com/?foo=bar&baz=foo&%F0%9D%95%A6=%F0%9D%95%A6") + assert url.extend_query("𝕦=𝕦") == expected + + +def test_extend_query_with_non_ascii_same_key(): + url = URL("http://example.com/?foo=bar&baz=foo&%F0%9D%95%A6=%F0%9D%95%A6") + expected = URL( + "http://example.com/?foo=bar&baz=foo" + "&%F0%9D%95%A6=%F0%9D%95%A6&%F0%9D%95%A6=%F0%9D%95%A6" + ) + assert url.extend_query({"𝕦": "𝕦"}) == expected diff --git a/yarl/_url.py b/yarl/_url.py index 658daf6b9..b73562259 100644 --- a/yarl/_url.py +++ b/yarl/_url.py @@ -1250,6 +1250,36 @@ def with_query(self, *args: Any, **kwargs: Any) -> "URL": new_query = self._get_str_query(*args, **kwargs) or "" return URL(self._val._replace(query=new_query), encoded=True) + @overload + def extend_query(self, query: Query) -> "URL": ... + + @overload + def extend_query(self, **kwargs: QueryVariable) -> "URL": ... + + def extend_query(self, *args: Any, **kwargs: Any) -> "URL": + """Return a new URL with query part combined with the existing. + + This method will not remove existing query parameters. + + Example: + >>> url = URL('http://example.com/?a=1&b=2') + >>> url.extend_query(a=3, c=4) + URL('http://example.com/?a=1&b=2&a=3&c=4') + """ + new_query_string = self._get_str_query(*args, **kwargs) + if not new_query_string: + return self + if current_query := self.raw_query_string: + # both strings are already encoded so we can use a simple + # string join + if current_query[-1] == "&": + combined_query = f"{current_query}{new_query_string}" + else: + combined_query = f"{current_query}&{new_query_string}" + else: + combined_query = new_query_string + return URL(self._val._replace(query=combined_query), encoded=True) + @overload def update_query(self, query: Query) -> "URL": ... @@ -1257,7 +1287,15 @@ def update_query(self, query: Query) -> "URL": ... def update_query(self, **kwargs: QueryVariable) -> "URL": ... def update_query(self, *args: Any, **kwargs: Any) -> "URL": - """Return a new URL with query part updated.""" + """Return a new URL with query part updated. + + This method will overwrite existing query parameters. + + Example: + >>> url = URL('http://example.com/?a=1&b=2') + >>> url.update_query(a=3, c=4) + URL('http://example.com/?a=3&b=2&c=4') + """ s = self._get_str_query(*args, **kwargs) if s is None: return URL(self._val._replace(query=""), encoded=True)