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

@@ -364,3 +364,60 @@ 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")

yarl/_url.py

@@ -29,7 +29,7 @@
 )
 
 import idna
-from multidict import MultiDict, MultiDictProxy
+from multidict import MultiDict, MultiDictProxy, istr
 
 from ._helpers import cached_property
 from ._quoting import _Quoter, _Unquoter
@@ -88,6 +88,7 @@ class _InternalURLCache(TypedDict, total=False):
     explicit_port: Union[int, None]
     raw_path: str
     path: str
+    _parsed_query: List[Tuple[str, str]]
     query: "MultiDictProxy[str]"
     raw_query_string: str
     query_string: str
@@ -697,6 +698,11 @@ def path(self) -> str:
         """
         return self._PATH_UNQUOTER(self.raw_path)
 
+    @cached_property
+    def _parsed_query(self) -> List[Tuple[str, str]]:
+        """Parse query part of URL."""
+        return parse_qsl(self.raw_query_string, keep_blank_values=True)
+
     @cached_property
     def query(self) -> "MultiDictProxy[str]":
         """A MultiDictProxy representing parsed query parameters in decoded
@@ -705,8 +711,7 @@ def query(self) -> "MultiDictProxy[str]":
         Empty value if URL has no query part.
 
         """
-        ret = MultiDict(parse_qsl(self.raw_query_string, keep_blank_values=True))
-        return MultiDictProxy(ret)
+        return MultiDictProxy(MultiDict(self._parsed_query))
 
     @cached_property
     def raw_query_string(self) -> str:
@@ -1154,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))
@@ -1176,8 +1181,17 @@ def _query_var(v: QueryVariable) -> str:
             "of type {}".format(v, cls)
         )
 
+    def _get_str_query_from_iterable(
+        self, items: "Iterable[Tuple[Union[str, istr], str]]"
+    ) -> str:
+        """Return a query string from an iterable."""
+        quoter = self._QUERY_PART_QUOTER
+        # A listcomp is used since listcomps are inlined on CPython 3.12+ and
+        # they are a bit faster than a generator expression.
+        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(
@@ -1190,32 +1204,29 @@ 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, (MultiDict, MultiDictProxy)):
+            return self._get_str_query_from_iterable(query.items())
+        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):
-            quoter = self._QUERY_PART_QUOTER
+        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`.
-            query = "&".join(
-                quoter(k) + "=" + quoter(self._query_var(v)) for k, v in query
-            )
-        else:
-            raise TypeError(
-                "Invalid query type: only str, mapping or "
-                "sequence of (key, value) pairs is allowed"
-            )
+            return self._get_str_query_from_iterable(query)
 
-        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": ...
@@ -1239,9 +1250,28 @@ def with_query(self, *args: Any, **kwargs: Any) -> "URL":
         # N.B. doesn't cleanup query/fragment
 
         new_query = self._get_str_query(*args, **kwargs) or ""
-        return URL(
-            self._val._replace(path=self._val.path, query=new_query), encoded=True
-        )
+        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 new_query_string is None:
+            return self
+        return self._merge_query(new_query_string, update=False)
 
     @overload
     def update_query(self, query: Query) -> "URL": ...
@@ -1250,17 +1280,35 @@ 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."""
-        s = self._get_str_query(*args, **kwargs)
-        query = None
-        if s is not None:
-            new_query = MultiDict(parse_qsl(s, keep_blank_values=True))
-            query = MultiDict(self.query)
-            query.update(new_query)
+        """Return a new URL with query part updated.
 
-        return URL(
-            self._val._replace(query=self._get_str_query(query) or ""), encoded=True
-        )
+        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')
+        """
+        new_query_string = self._get_str_query(*args, **kwargs)
+        if new_query_string is None:
+            return URL(self._val._replace(query=""), encoded=True)
+        return self._merge_query(new_query_string, update=True)
+
+    def _merge_query(self, new_query_string: Union[str, None], update: bool) -> "URL":
+        """Return a new URL with query part merged or extended."""
+        new_parsed = parse_qsl(new_query_string, keep_blank_values=True)
+
+        if update:
+            new_query = MultiDict(self._parsed_query)
+            new_query.update(new_parsed)
+        else:
+            new_query = MultiDict(self._parsed_query + new_parsed)
+
+        # We can use the faster _get_str_query_from_iterable here because
+        # we constructed the MultiDict ourselves and we know there are
+        # no QueryVariable as the values in this case.
+        combined_query = self._get_str_query_from_iterable(new_query.items()) or ""
+        return URL(self._val._replace(query=combined_query), encoded=True)
 
     def without_query_params(self, *query_params: str) -> "URL":
         """Remove some keys from query part and return new URL."""