Add extend_query method (#1128)

CHANGES/1128.feature.rst

@@ -0,0 +1,3 @@
+Added :meth:`URL.extend_query() <yarl.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``.

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)
 

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

yarl/_url.py

@@ -1250,14 +1250,52 @@ 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": ...
 
     @overload
     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)