Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 48 additions & 7 deletions src/osw/wiki_tools.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import getpass
import warnings
from typing import Dict, List, Optional, Tuple, Union

import mwclient
Expand Down Expand Up @@ -216,6 +217,33 @@ def prefix_search_(single_text) -> Union[List[str], dict]:
# return page_list # original return


def _ask_results_as_dict(results: Union[dict, list]) -> dict:
"""Normalise the ``results`` payload of an SMW ``ask`` response to a mapping.

SMW serialises a non-empty result set as a JSON object keyed by page title,
but an empty one as a JSON array, which would otherwise break attribute
access on the result.

Parameters
----------
results :
The ``result["query"]["results"]`` payload of an SMW ``ask`` response.

Returns
-------
result:
The payload normalised to a dict, keyed by page title (or index if a
title is unavailable).
"""
if isinstance(results, dict):
return results
if not results:
return {}
return {
page.get("fulltext", str(index)): page for index, page in enumerate(results)
}


def semantic_search(
site: mwclient.client.Site, query: Union[str, List[str], SearchParam]
) -> Union[List[str], List[dict]]:
Expand All @@ -242,26 +270,39 @@ def semantic_search_(single_query):
page_list = list()
single_query += f"|limit={query.limit}"
result = site.api("ask", query=single_query, format="json")
results = _ask_results_as_dict(result["query"]["results"])
n = len(results)
if query.debug:
if len(result["query"]["results"]) == 0:
if n == 0:
print(f"Query '{single_query}' returned no results")
else:
print(
"Query '{}' returned {} results".format(
single_query, len(result["query"]["results"])
)
)
print(f"Query '{single_query}' returned {n} results")
if n >= query.limit:
warnings.warn(
f"Query '{single_query}' returned {n} results, which meets the "
f"requested limit of {query.limit}. Results are truncated - raise "
f"the limit or page through with '|offset=' to retrieve the "
f"remainder."
)
if query.return_json:
return result

for page in result["query"]["results"].values():
dropped = 0
for page in results.values():
title = page["fulltext"]
exists = page["exists"]
if "#" not in title and query.debug:
print(title)
# original position of "page_list.append(title)" line
if exists == "1":
page_list.append(title)
else:
dropped += 1
if dropped > 0:
warnings.warn(
f"Query '{single_query}': {dropped} of {n} results were dropped "
f"because the wiki reported them as non-existing pages."
)
return page_list

if query.parallel:
Expand Down
129 changes: 129 additions & 0 deletions tests/test_wiki_tools.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,8 @@
import warnings
from unittest.mock import MagicMock

import pytest

import osw.wiki_tools as wt


Expand Down Expand Up @@ -29,6 +32,12 @@ def _ask_result(*titles):
}


def _ask_result_empty():
"""Build the SMW ``ask`` API result dict for a zero-result query, mirroring
SMW's behaviour of serialising an empty result set as a JSON array."""
return {"query": {"results": []}}


def test_semantic_search_return_json_single_query_returns_list_with_dict():
result = _ask_result("Item:OSW1")
site = MagicMock()
Expand Down Expand Up @@ -78,6 +87,126 @@ def test_semantic_search_returns_flat_list_of_titles():
assert out == ["Item:OSW1", "Item:OSW2"]


def test_semantic_search_zero_results_returns_empty_list():
site = MagicMock()
site.api.return_value = _ask_result_empty()

out = wt.semantic_search(site, "[[HasType::Category:Nonexistent]]")

assert out == []


def test_semantic_search_zero_results_return_json_returns_list_with_dict():
result = _ask_result_empty()
site = MagicMock()
site.api.return_value = result

out = wt.semantic_search(
site,
wt.SearchParam(query="[[HasType::Category:Nonexistent]]", return_json=True),
)

assert out == [result]


def test_semantic_search_batch_with_one_zero_result_query():
result_a = _ask_result("Item:OSW1")
result_b = _ask_result_empty()
result_c = _ask_result("Item:OSW3")
site = MagicMock()
site.api.side_effect = [result_a, result_b, result_c]

out = wt.semantic_search(
site,
wt.SearchParam(
query=[
"[[HasType::Category:A]]",
"[[HasType::Category:B]]",
"[[HasType::Category:C]]",
]
),
)

assert out == ["Item:OSW1", "Item:OSW3"]


def test_semantic_search_parallel_batch_with_one_zero_result_query():
results = [
_ask_result("Item:OSW1"),
_ask_result("Item:OSW2"),
_ask_result_empty(),
_ask_result("Item:OSW4"),
_ask_result("Item:OSW5"),
_ask_result("Item:OSW6"),
]
site = MagicMock()
site.api.side_effect = results

out = wt.semantic_search(
site,
wt.SearchParam(
query=[
"[[HasType::Category:A]]",
"[[HasType::Category:B]]",
"[[HasType::Category:C]]",
"[[HasType::Category:D]]",
"[[HasType::Category:E]]",
"[[HasType::Category:F]]",
]
),
)

assert sorted(out) == [
"Item:OSW1",
"Item:OSW2",
"Item:OSW4",
"Item:OSW5",
"Item:OSW6",
]


def test_semantic_search_truncation_warning():
titles = [f"Item:OSW{i}" for i in range(5)]
result = _ask_result(*titles)
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=5)
)

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)
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]]", limit=1000)
)

assert not any("truncated" in str(w.message) for w in caught)
assert sorted(out) == sorted(titles)


def test_semantic_search_exists_drop_warning():
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, "[[HasType::Category:Item]]")

assert out == ["Item:OSW1"]


def _prefixsearch_result(*titles):
"""Build a minimal MediaWiki ``prefixsearch`` API result dict."""
return {
Expand Down
Loading