fix(search): keep a limit the caller wrote into the ask query

Merge pull request #161 from OpenSemanticLab/fix/semantic-search-respect-query-limit

Author: LukasGold

src/osw/wiki_tools.py

@@ -1,4 +1,5 @@
 import getpass
+import re
 import warnings
 from typing import Dict, List, Optional, Tuple, Union
 
@@ -137,6 +138,9 @@ class SearchParam(OswBaseModel):
     parallel: Optional[bool] = None  # is set to true if query is a list longer than 5
     debug: Optional[bool] = False
     limit: Optional[int] = 1000
+    """the result limit to apply, None to apply none and leave it to the wiki.
+    Ignored by semantic_search for a query that sets 'limit=' itself, since SMW
+    honours the last limit in the query string"""
     return_json: Optional[bool] = False
 
     def __init__(self, **data):
@@ -244,6 +248,36 @@ def _ask_results_as_dict(results: Union[dict, list]) -> dict:
     }
 
 
+CONDITION_PATTERN = re.compile(r"\[\[.*?\]\]", re.DOTALL)
+"""matches an SMW query condition, which may contain '|' as the disjunction
+operator and must therefore be removed before the parameters are split"""
+LIMIT_PARAM_PATTERN = re.compile(r"^limit\s*=\s*(\d+)$", re.IGNORECASE)
+"""matches a limit parameter of an SMW query, e.g. 'limit=100'"""
+
+
+def get_query_limit(query: str) -> Optional[int]:
+    """Returns the limit an SMW ask query sets itself, None if it sets none
+
+    Parameters
+    ----------
+    query :
+        an SMW ask query string, e.g. '[[Category:Item]]|?Name|limit=2'
+
+    Returns
+    -------
+    result:
+        the limit the query asks for, or None if the query does not set one or
+        sets one that is not a number
+    """
+    limit = None
+    for parameter in CONDITION_PATTERN.sub("", query).split("|"):
+        match = LIMIT_PARAM_PATTERN.match(parameter.strip())
+        if match:
+            # SMW honours the last limit in the query string
+            limit = int(match.group(1))
+    return limit
+
+
 def semantic_search(
     site: mwclient.client.Site, query: Union[str, List[str], SearchParam]
 ) -> Union[List[str], List[dict]]:
@@ -254,7 +288,9 @@ def semantic_search(
     site :
         Site object from mwclient lib
     query :
-        (List of) query text(s) or instance of SearchParam
+        (List of) query text(s) or instance of SearchParam. A query that sets
+        ``limit=`` itself keeps that limit; ``SearchParam.limit`` is only
+        appended to a query that does not, and only when it is not None.
 
     Returns
     -------
@@ -268,7 +304,14 @@ def semantic_search(
 
     def semantic_search_(single_query):
         page_list = list()
-        single_query += f"|limit={query.limit}"
+        # SMW honours the last limit in the query string, so appending the
+        # default would silently override a limit the caller wrote themselves.
+        # A SearchParam limit of None asks for no limit at all, leaving the
+        # wiki to apply its own default
+        limit = get_query_limit(single_query)
+        if limit is None and query.limit is not None:
+            limit = query.limit
+            single_query += f"|limit={limit}"
         result = site.api("ask", query=single_query, format="json")
         results = _ask_results_as_dict(result["query"]["results"])
         n = len(results)
@@ -277,10 +320,13 @@ def semantic_search_(single_query):
                 print(f"Query '{single_query}' returned no results")
             else:
                 print(f"Query '{single_query}' returned {n} results")
-        if n >= query.limit:
+        # No limit in force, or 'limit=0' asking for no results at all as a
+        # count format does, means the result count says nothing about
+        # truncation
+        if limit and n >= limit:
             warnings.warn(
                 f"Query '{single_query}' returned {n} results, which meets the "
-                f"requested limit of {query.limit}. Results are truncated - raise "
+                f"requested limit of {limit}. Results are truncated - raise "
                 f"the limit or page through with '|offset=' to retrieve the "
                 f"remainder."
             )

tests/test_wiki_tools.py

@@ -207,6 +207,138 @@ def test_semantic_search_exists_drop_warning():
     assert out == ["Item:OSW1"]
 
 
+@pytest.mark.parametrize(
+    "query, expected",
+    [
+        ("[[HasType::Category:Item]]", None),
+        ("[[HasType::Category:Item]]|limit=2", 2),
+        ("[[HasType::Category:Item]]|?Name|limit=2|offset=5", 2),
+        ("[[HasType::Category:Item]]| limit = 2 ", 2),
+        ("[[HasType::Category:Item]]|Limit=2", 2),
+        # SMW honours the last one, so that is the effective limit
+        ("[[HasType::Category:Item]]|limit=2|limit=7", 7),
+        # a limit inside a condition is a value, not a parameter
+        ("[[HasText::limit=2]]", None),
+        ("[[HasText::a||limit=2]]", None),
+        # a printout named 'limit', not the parameter
+        ("[[HasType::Category:Item]]|?limit", None),
+        # a printout parameter, which limits that printout and not the query
+        ("[[HasType::Category:Item]]|?Has subobject|+limit=3", None),
+        # not a number, so there is no limit to honour
+        ("[[HasType::Category:Item]]|limit=all", None),
+    ],
+)
+def test_get_query_limit(query, expected):
+    assert wt.get_query_limit(query) == expected
+
+
+def test_semantic_search_appends_the_default_limit():
+    site = MagicMock()
+    site.api.return_value = _ask_result("Item:OSW1")
+
+    wt.semantic_search(site, "[[HasType::Category:Item]]")
+
+    assert site.api.call_args.kwargs["query"] == (
+        "[[HasType::Category:Item]]|limit=1000"
+    )
+
+
+def test_semantic_search_keeps_a_limit_the_caller_wrote():
+    """Appending the default would override it, since SMW honours the last one."""
+    site = MagicMock()
+    site.api.return_value = _ask_result("Item:OSW1")
+
+    wt.semantic_search(site, "[[HasType::Category:Item]]|limit=2")
+
+    assert site.api.call_args.kwargs["query"] == "[[HasType::Category:Item]]|limit=2"
+
+
+def test_semantic_search_query_limit_beats_the_search_param_limit():
+    site = MagicMock()
+    site.api.return_value = _ask_result("Item:OSW1")
+
+    wt.semantic_search(
+        site, wt.SearchParam(query="[[HasType::Category:Item]]|limit=2", limit=500)
+    )
+
+    assert site.api.call_args.kwargs["query"] == "[[HasType::Category:Item]]|limit=2"
+
+
+def test_semantic_search_truncation_warning_uses_the_query_limit():
+    """The caller's limit is the one the results were truncated at."""
+    titles = [f"Item:OSW{i}" for i in range(2)]
+    site = MagicMock()
+    site.api.return_value = _ask_result(*titles)
+
+    with pytest.warns(UserWarning, match="requested limit of 2"):
+        out = wt.semantic_search(site, "[[HasType::Category:Item]]|limit=2")
+
+    assert sorted(out) == sorted(titles)
+
+
+def test_semantic_search_limit_none_sends_no_limit():
+    """None asks for no limit at all, leaving the wiki to apply its own."""
+    site = MagicMock()
+    site.api.return_value = _ask_result("Item:OSW1")
+
+    wt.semantic_search(
+        site, wt.SearchParam(query="[[HasType::Category:Item]]", limit=None)
+    )
+
+    assert site.api.call_args.kwargs["query"] == "[[HasType::Category:Item]]"
+
+
+def test_semantic_search_limit_none_keeps_a_limit_the_caller_wrote():
+    site = MagicMock()
+    site.api.return_value = _ask_result("Item:OSW1")
+
+    wt.semantic_search(
+        site, wt.SearchParam(query="[[HasType::Category:Item]]|limit=2", limit=None)
+    )
+
+    assert site.api.call_args.kwargs["query"] == "[[HasType::Category:Item]]|limit=2"
+
+
+def test_semantic_search_limit_none_does_not_warn_about_truncation():
+    """With no limit in force the result count says nothing about truncation."""
+    titles = [f"Item:OSW{i}" for i in range(5)]
+    site = MagicMock()
+    site.api.return_value = _ask_result(*titles)
+
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        out = wt.semantic_search(
+            site, wt.SearchParam(query="[[HasType::Category:Item]]", limit=None)
+        )
+
+    assert not any("truncated" in str(w.message) for w in caught)
+    assert sorted(out) == sorted(titles)
+
+
+def test_semantic_search_no_truncation_warning_for_a_zero_limit():
+    """'limit=0' asks for no results, so meeting it is not truncation."""
+    site = MagicMock()
+    site.api.return_value = _ask_result_empty()
+
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        wt.semantic_search(site, "[[HasType::Category:Item]]|limit=0")
+
+    assert not any("truncated" in str(w.message) for w in caught)
+
+
+def test_semantic_search_no_truncation_warning_below_the_query_limit():
+    titles = [f"Item:OSW{i}" for i in range(2)]
+    site = MagicMock()
+    site.api.return_value = _ask_result(*titles)
+
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        wt.semantic_search(site, "[[HasType::Category:Item]]|limit=50")
+
+    assert not any("truncated" in str(w.message) for w in caught)
+
+
 def _prefixsearch_result(*titles):
     """Build a minimal MediaWiki ``prefixsearch`` API result dict."""
     return {