Add extend_query method

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` the 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

@@ -1159,19 +1159,19 @@ def _query_seq_pairs(
     @staticmethod
     def _query_var(v: QueryVariable) -> str:
         cls = type(v)
-        if issubclass(cls, str):
+        if cls is str or issubclass(cls, str):
             if TYPE_CHECKING:
                 assert isinstance(v, str)
             return v
-        if issubclass(cls, float):
+        if cls is float or issubclass(cls, float):
             if TYPE_CHECKING:
                 assert isinstance(v, float)
             if math.isinf(v):
                 raise ValueError("float('inf') is not supported")
             if math.isnan(v):
                 raise ValueError("float('nan') is not supported")
             return str(float(v))
-        if issubclass(cls, int) and cls is not bool:
+        if cls is int or (issubclass(cls, int) and cls is not bool):
             if TYPE_CHECKING:
                 assert isinstance(v, int)
             return str(int(v))
@@ -1191,7 +1191,7 @@ def _get_str_query_from_iterable(
         return "&".join([f"{quoter(k)}={quoter(self._query_var(v))}" for k, v in items])
 
     def _get_str_query(self, *args: Any, **kwargs: Any) -> Union[str, None]:
-        query: Union[str, Mapping[str, QueryVariable], None]
+        query: Union[str, "Mapping[str, QueryVariable]", None]
         if kwargs:
             if len(args) > 0:
                 raise ValueError(
@@ -1204,29 +1204,27 @@ def _get_str_query(self, *args: Any, **kwargs: Any) -> Union[str, None]:
             raise ValueError("Either kwargs or single query parameter must be present")
 
         if query is None:
-            query = None
-        elif isinstance(query, Mapping):
+            return None
+        if isinstance(query, Mapping):
             quoter = self._QUERY_PART_QUOTER
-            query = "&".join(self._query_seq_pairs(quoter, query.items()))
-        elif isinstance(query, str):
-            query = self._QUERY_QUOTER(query)
-        elif isinstance(query, (bytes, bytearray, memoryview)):
+            return "&".join(self._query_seq_pairs(quoter, query.items()))
+        if isinstance(query, str):
+            return self._QUERY_QUOTER(query)
+        if isinstance(query, (bytes, bytearray, memoryview)):
             raise TypeError(
                 "Invalid query type: bytes, bytearray and memoryview are forbidden"
             )
-        elif isinstance(query, Sequence):
+        if isinstance(query, Sequence):
             # We don't expect sequence values if we're given a list of pairs
             # already; only mappings like builtin `dict` which can't have the
             # same key pointing to multiple values are allowed to use
             # `_query_seq_pairs`.
             return self._get_str_query_from_iterable(query)
-        else:
-            raise TypeError(
-                "Invalid query type: only str, mapping or "
-                "sequence of (key, value) pairs is allowed"
-            )
 
-        return query
+        raise TypeError(
+            "Invalid query type: only str, mapping or "
+            "sequence of (key, value) pairs is allowed"
+        )
 
     @overload
     def with_query(self, query: Query) -> "URL": ...
@@ -1252,14 +1250,50 @@ 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:
+            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)