diff --git a/src/osw/wiki_tools.py b/src/osw/wiki_tools.py index 7139ee2..2cef2cb 100644 --- a/src/osw/wiki_tools.py +++ b/src/osw/wiki_tools.py @@ -135,6 +135,26 @@ def create_site_object( return site +class SemanticSearchResult(OswBaseModel): + """Result of a single semantic query, including whether the wiki truncated it""" + + query: str + """the query as sent to the wiki, including any limit appended by + semantic_search""" + titles: List[str] + """the page-title fulltext strings of the results that exist""" + count: int + """the number of results this response carried, before dropping + non-existing pages, so it can be larger than len(titles). It is not the + total number of matching pages on the wiki, which SMW does not report""" + truncated: bool + """True if the wiki reported further results beyond those returned""" + next_offset: Optional[int] = None + """the absolute offset at which the remaining results start, to be passed + back as '|offset='. None for a complete result set. It says where to + continue, not how many results remain""" + + class SearchParam(OswBaseModel): """Search parameters for semantic and prefix search""" @@ -146,6 +166,11 @@ class SearchParam(OswBaseModel): 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 + return_meta: Optional[bool] = False + """If True, semantic_search returns one SemanticSearchResult per query instead + of a flat list of titles, so that a caller can report a truncated result set. + Ignored when return_json is True, since the raw wiki response already carries + the truncation signal""" def __init__(self, **data): super().__init__(**data) @@ -284,7 +309,7 @@ def get_query_limit(query: str) -> Optional[int]: def semantic_search( site: mwclient.client.Site, query: Union[str, List[str], SearchParam] -) -> Union[List[str], List[dict]]: +) -> Union[List[str], List[dict], List[SemanticSearchResult]]: """Semantic query Parameters @@ -299,9 +324,12 @@ def semantic_search( Returns ------- result: - With ``return_json=False`` (default): a flat list of page-title fulltext - strings. With ``return_json=True``: a list of raw SMW ``ask`` result dicts, - one per query (always a list, even for a single query). + With ``return_json=False`` and ``return_meta=False`` (default): a flat + list of page-title fulltext strings. With ``return_json=True``: a list of + raw SMW ``ask`` result dicts, one per query (always a list, even for a + single query). With ``return_meta=True``: a list of SemanticSearchResult, + one per query, which reports whether the wiki truncated the result set. + ``return_json`` takes precedence if both are set. """ if not isinstance(query, SearchParam): query = SearchParam(query=query) @@ -324,15 +352,21 @@ def semantic_search_(single_query): print(f"Query '{single_query}' returned no results") else: print(f"Query '{single_query}' returned {n} results") - # 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: + # SMW reports an incomplete result set with a top-level + # 'query-continue-offset' holding the offset the remainder starts at, + # and omits the key for a complete one. That replaces the earlier + # comparison of the result count against the limit, which was wrong in + # both directions: it could not see the wiki's own '$smwgQMaxLimit' + # cap, and it reported a complete set of exactly 'limit' results as + # truncated + next_offset = result.get("query-continue-offset") + truncated = next_offset is not None + if truncated: warnings.warn( - f"Query '{single_query}' returned {n} results, which meets the " - f"requested limit of {limit}. Results are truncated - raise " - f"the limit or page through with '|offset=' to retrieve the " - f"remainder." + f"Query '{single_query}' returned {n} results and the wiki " + f"reports further ones. Results are truncated - raise the " + f"limit or page through with '|offset={next_offset}' to " + f"retrieve the remainder." ) if query.return_json: return result @@ -353,6 +387,14 @@ def semantic_search_(single_query): f"Query '{single_query}': {dropped} of {n} results were dropped " f"because the wiki reported them as non-existing pages." ) + if query.return_meta: + return SemanticSearchResult( + query=single_query, + titles=page_list, + count=n, + truncated=truncated, + next_offset=next_offset, + ) return page_list if query.parallel: @@ -362,10 +404,10 @@ def semantic_search_(single_query): else: query_results = [semantic_search_(single_query=sq) for sq in query.query] - if query.return_json: - # Each entry of query_results is the raw SMW result dict for one query. - # Do not flatten dicts; always return the list of result dicts (one per - # query), even when only a single query was passed. + if query.return_json or query.return_meta: + # Each entry of query_results is the raw SMW result dict, or the + # SemanticSearchResult, for one query. Do not flatten those; always + # return one entry per query, even when only a single query was passed. return query_results return [item for sublist in query_results for item in sublist] diff --git a/src/osw/wtsite.py b/src/osw/wtsite.py index 8193c3e..5dc8ce9 100644 --- a/src/osw/wtsite.py +++ b/src/osw/wtsite.py @@ -568,7 +568,9 @@ def semantic_search(self, query: Union[str, SearchParam]): Returns ------- - A list of page titles + A list of page titles, or, if the SearchParam sets return_json or + return_meta, one raw result dict or SemanticSearchResult per query. + See wiki_tools.semantic_search for details. """ return wt.semantic_search(self._site, query) @@ -621,6 +623,12 @@ def modify_search_results( dryrun Deprecated, use param.dryrun instead. if True, no actual changes are made, by default False + + Raises + ------ + ValueError + If param.query is a SearchParam asking for anything other than page + titles, meaning return_json or return_meta. """ if not isinstance(param, WtSite.ModifySearchResultsParam): param = WtSite.ModifySearchResultsParam( @@ -631,6 +639,23 @@ def modify_search_results( dryrun=dryrun, ) + # Both searches can be asked for the raw API response, and the semantic + # one for SemanticSearchResult objects. This method looks up and edits + # a page per result, so it needs the titles themselves. Reject the + # other options here rather than failing further down on a dict where + # a title is expected + if isinstance(param.query, wt.SearchParam): + unsupported = [ + name + for name in ("return_json", "return_meta") + if getattr(param.query, name) + ] + if unsupported: + raise ValueError( + f"modify_search_results edits a page per search result, so " + f"the query must not set {' or '.join(unsupported)}." + ) + titles = [] if param.mode == "prefix": titles = wt.prefix_search(self._site, param.query) diff --git a/tests/test_wiki_tools.py b/tests/test_wiki_tools.py index f1c2e4f..d1403e3 100644 --- a/tests/test_wiki_tools.py +++ b/tests/test_wiki_tools.py @@ -48,9 +48,12 @@ def test_read_domains_from_credentials_file_valid_file_returns_domains_and_accou assert accounts == {"example.org": {"username": "user", "password": "pass"}} -def _ask_result(*titles): - """Build a minimal SMW ``ask`` API result dict for the given page titles.""" - return { +def _ask_result(*titles, continue_offset=None): + """Build a minimal SMW ``ask`` API result dict for the given page titles. + + ``continue_offset`` adds the top-level ``query-continue-offset`` key that SMW + sends when further results exist beyond the ones returned.""" + result = { "query": { "results": { title: { @@ -65,6 +68,9 @@ def _ask_result(*titles): } } } + if continue_offset is not None: + result["query-continue-offset"] = continue_offset + return result def _ask_result_empty(): @@ -202,7 +208,7 @@ def test_semantic_search_parallel_batch_with_one_zero_result_query(): def test_semantic_search_truncation_warning(): titles = [f"Item:OSW{i}" for i in range(5)] - result = _ask_result(*titles) + result = _ask_result(*titles, continue_offset=5) site = MagicMock() site.api.return_value = result @@ -214,6 +220,23 @@ def test_semantic_search_truncation_warning(): assert sorted(out) == sorted(titles) +def test_semantic_search_no_truncation_warning_for_a_complete_set_at_the_limit(): + """A wiki that has exactly as many results as the limit asked for sends no + offset, so meeting the limit is not on its own a sign of 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=5) + ) + + assert not any("truncated" in str(w.message) for w in caught) + assert sorted(out) == sorted(titles) + + def test_semantic_search_no_truncation_warning_below_limit(): titles = [f"Item:OSW{i}" for i in range(5)] result = _ask_result(*titles) @@ -242,6 +265,160 @@ def test_semantic_search_exists_drop_warning(): assert out == ["Item:OSW1"] +def test_semantic_search_truncation_warning_from_the_continue_offset(): + # Fewer results than the limit, so only the wiki's own signal can report + # that the result set was cut short + result = _ask_result("Item:OSW1", "Item:OSW2", continue_offset=2) + site = MagicMock() + site.api.return_value = result + + with pytest.warns(UserWarning, match="truncated"): + out = wt.semantic_search( + site, wt.SearchParam(query="[[HasType::Category:Item]]", limit=1000) + ) + + assert sorted(out) == ["Item:OSW1", "Item:OSW2"] + + +def test_semantic_search_limit_none_warns_from_the_continue_offset(): + # With no limit in force the wiki applies its own '$smwgQMaxLimit' cap, + # which the result count cannot detect + result = _ask_result("Item:OSW1", "Item:OSW2", continue_offset=2) + site = MagicMock() + site.api.return_value = result + + with pytest.warns(UserWarning, match="truncated"): + wt.semantic_search( + site, wt.SearchParam(query="[[HasType::Category:Item]]", limit=None) + ) + + +def test_semantic_search_return_json_keeps_the_continue_offset(): + result = _ask_result("Item:OSW1", continue_offset=1) + site = MagicMock() + site.api.return_value = result + + with pytest.warns(UserWarning, match="truncated"): + out = wt.semantic_search( + site, wt.SearchParam(query="[[HasType::Category:Item]]", return_json=True) + ) + + # The raw response is returned unchanged, so it already carries the + # truncation signal + assert out[0]["query-continue-offset"] == 1 + + +def test_semantic_search_return_meta_reports_truncation(): + result = _ask_result("Item:OSW1", "Item:OSW2", continue_offset=2) + site = MagicMock() + site.api.return_value = result + + with pytest.warns(UserWarning, match="truncated"): + out = wt.semantic_search( + site, wt.SearchParam(query="[[HasType::Category:Item]]", return_meta=True) + ) + + assert len(out) == 1 + assert isinstance(out[0], wt.SemanticSearchResult) + assert out[0].truncated is True + assert out[0].next_offset == 2 + assert out[0].count == 2 + assert sorted(out[0].titles) == ["Item:OSW1", "Item:OSW2"] + # the limit appended by semantic_search is part of the reported query + assert out[0].query == "[[HasType::Category:Item]]|limit=1000" + + +def test_semantic_search_return_meta_reports_a_complete_result(): + result = _ask_result("Item:OSW1", "Item:OSW2") + site = MagicMock() + site.api.return_value = result + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + out = wt.semantic_search( + site, wt.SearchParam(query="[[HasType::Category:Item]]", return_meta=True) + ) + + assert not any("truncated" in str(w.message) for w in caught) + assert out[0].truncated is False + assert out[0].next_offset is None + assert out[0].count == 2 + + +def test_semantic_search_return_meta_returns_one_result_per_query(): + site = MagicMock() + site.api.side_effect = [ + _ask_result("Item:OSW1", continue_offset=1), + _ask_result("Item:OSW2"), + ] + + with pytest.warns(UserWarning, match="truncated"): + out = wt.semantic_search( + site, + wt.SearchParam( + query=["[[HasType::Category:A]]", "[[HasType::Category:B]]"], + return_meta=True, + ), + ) + + # One entry per query, not flattened into a single list of titles + assert len(out) == 2 + assert [r.truncated for r in out] == [True, False] + assert [r.titles for r in out] == [["Item:OSW1"], ["Item:OSW2"]] + + +def test_semantic_search_return_meta_count_includes_dropped_pages(): + result = _ask_result("Item:OSW1", "Item:OSW2") + result["query"]["results"]["Item:OSW2"]["exists"] = "" + site = MagicMock() + site.api.return_value = result + + with pytest.warns(UserWarning, match="non-existing"): + out = wt.semantic_search( + site, wt.SearchParam(query="[[HasType::Category:Item]]", return_meta=True) + ) + + assert out[0].count == 2 + assert out[0].titles == ["Item:OSW1"] + + +def test_semantic_search_return_meta_keeps_the_query_order_when_parallel(): + # More than five queries makes SearchParam switch to the parallel path, where + # the results must still line up with the queries they came from + queries = [f"[[HasType::Category:C{i}]]" for i in range(6)] + + def api(action, query=None, format=None): + index = queries.index(query.split("|limit=")[0]) + return _ask_result(f"Item:OSW{index}", continue_offset=index or None) + + site = MagicMock() + site.api.side_effect = api + param = wt.SearchParam(query=queries, return_meta=True) + assert param.parallel is True + + with pytest.warns(UserWarning, match="truncated"): + out = wt.semantic_search(site, param) + + assert [r.query.split("|limit=")[0] for r in out] == queries + assert [r.titles for r in out] == [[f"Item:OSW{i}"] for i in range(6)] + assert [r.truncated for r in out] == [False] + [True] * 5 + + +def test_semantic_search_return_json_takes_precedence_over_return_meta(): + result = _ask_result("Item:OSW1") + site = MagicMock() + site.api.return_value = result + + out = wt.semantic_search( + site, + wt.SearchParam( + query="[[HasType::Category:Item]]", return_json=True, return_meta=True + ), + ) + + assert out == [result] + + @pytest.mark.parametrize( "query, expected", [ @@ -299,13 +476,14 @@ def test_semantic_search_query_limit_beats_the_search_param_limit(): 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.""" +def test_semantic_search_truncation_warning_names_the_query_as_sent(): + """The caller's limit is the one the results were truncated at, so the + warning has to quote the query carrying it.""" titles = [f"Item:OSW{i}" for i in range(2)] site = MagicMock() - site.api.return_value = _ask_result(*titles) + site.api.return_value = _ask_result(*titles, continue_offset=2) - with pytest.warns(UserWarning, match="requested limit of 2"): + with pytest.warns(UserWarning, match=r"\|limit=2'"): out = wt.semantic_search(site, "[[HasType::Category:Item]]|limit=2") assert sorted(out) == sorted(titles) @@ -334,8 +512,9 @@ def test_semantic_search_limit_none_keeps_a_limit_the_caller_wrote(): 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.""" +def test_semantic_search_limit_none_does_not_warn_without_an_offset(): + """A wiki reporting no further results never triggers the warning, whatever + limit was in force.""" titles = [f"Item:OSW{i}" for i in range(5)] site = MagicMock() site.api.return_value = _ask_result(*titles) @@ -351,7 +530,8 @@ def test_semantic_search_limit_none_does_not_warn_about_truncation(): def test_semantic_search_no_truncation_warning_for_a_zero_limit(): - """'limit=0' asks for no results, so meeting it is not truncation.""" + """'limit=0' asks for a count rather than results, and SMW sends no offset + for it even when the query does match pages.""" site = MagicMock() site.api.return_value = _ask_result_empty() diff --git a/tests/test_wtsite_modify_search_results.py b/tests/test_wtsite_modify_search_results.py index 3e86273..325e94b 100644 --- a/tests/test_wtsite_modify_search_results.py +++ b/tests/test_wtsite_modify_search_results.py @@ -7,6 +7,9 @@ import threading +import pytest + +import osw.wiki_tools as wt import osw.wtsite as wtsite_mod from osw.wtsite import WtPage, WtSite @@ -144,3 +147,48 @@ def test_modify_search_results_legacy_keyword_call_still_works(monkeypatch): assert handled == titles assert comments == [("Item:OSW1", "[bot] legacy")] + + +@pytest.mark.parametrize("mode", ["prefix", "semantic"]) +@pytest.mark.parametrize("flag", ["return_json", "return_meta"]) +def test_modify_search_results_rejects_a_query_not_asking_for_titles( + monkeypatch, mode, flag +): + """A search asked for raw responses or result objects yields no titles to + edit, so the method must say so instead of failing on the first result.""" + ws = _make_fake_wtsite() + # Would hand back dicts rather than titles if the guard let the call through + _stub_search(monkeypatch, [{"fulltext": "Item:OSW1"}]) + + param = WtSite.ModifySearchResultsParam( + mode=mode, + query=wt.SearchParam(query="[[Category:Item]]", **{flag: True}), + comment="[bot] test", + ) + + with pytest.raises(ValueError, match=flag): + ws.modify_search_results(param, modify_page=lambda wtpage: None) + + +@pytest.mark.parametrize("mode", ["prefix", "semantic"]) +def test_modify_search_results_accepts_a_plain_search_param(monkeypatch, mode): + """The guard must not reject a SearchParam that does ask for titles.""" + ws = _make_fake_wtsite() + titles = ["Item:OSW1"] + pages = _make_pages(ws, titles) + comments = [] + + _stub_search(monkeypatch, titles) + _stub_get_page(monkeypatch, ws, pages) + _stub_edits(monkeypatch, pages, comments) + + handled = [] + + param = WtSite.ModifySearchResultsParam( + mode=mode, + query=wt.SearchParam(query="[[Category:Item]]"), + comment="[bot] test", + ) + ws.modify_search_results(param, modify_page=lambda wtpage: handled.append(1)) + + assert handled == [1]