diff --git a/add_links/api.py b/add_links/api.py index d84f4d6..e0174b7 100644 --- a/add_links/api.py +++ b/add_links/api.py @@ -1,5 +1,5 @@ import re -from typing import Any +import typing import requests from requests.adapters import HTTPAdapter @@ -8,6 +8,8 @@ from simplejson.scanner import JSONDecodeError from .language import get_current_language from .util import is_disambig +StrDict = dict[str, typing.Any] + ua = ( "find-link/2.2 " + "(https://github.com/EdwardBetts/find_link; contact: edward@4angle.com)" @@ -20,42 +22,47 @@ def get_query_url() -> str: return f"https://{get_current_language()}.wikipedia.org/w/api.php" -sessions = {} +sessions: dict[str, requests.sessions.Session] = {} -def get_session(): +def get_session() -> requests.sessions.Session: + """Get requests session.""" lang = get_current_language() if lang in sessions: return sessions[lang] s = requests.Session() s.headers = {"User-Agent": ua} s.mount("https://en.wikipedia.org", HTTPAdapter(max_retries=10)) - s.params = { - "format": "json", - "action": "query", - "formatversion": 2, - } + s.params = typing.cast( + dict[str, str | int], + { + "format": "json", + "action": "query", + "formatversion": 2, + }, + ) sessions[lang] = s return s class MediawikiError(Exception): - pass + """Mediawiki error.""" class MultipleRedirects(Exception): - pass + """Multiple redirects.""" class IncompleteReply(Exception): - pass + """Incomplete reply.""" class MissingPage(Exception): - pass + """Missing page.""" -def check_for_error(json_data): +def check_for_error(json_data: dict[str, typing.Any]) -> None: + """Check MediaWiki API reply for error.""" if "error" in json_data: raise MediawikiError(json_data["error"]["info"]) @@ -65,13 +72,13 @@ webpage_error = ( ) -def api_get(params: dict[str, Any]) -> dict[str, Any]: +def api_get(params: StrDict) -> StrDict: """Make call to Wikipedia API.""" s = get_session() r = s.get(get_query_url(), params=params) try: - ret = r.json() + ret: StrDict = r.json() except JSONDecodeError: if webpage_error in r.text: raise MediawikiError(webpage_error) @@ -81,22 +88,23 @@ def api_get(params: dict[str, Any]) -> dict[str, Any]: return ret -def get_first_page(params: dict[str, str]) -> dict[str, Any]: +def get_first_page(params: dict[str, str]) -> StrDict: """Run Wikipedia API query and return the first page.""" - page = api_get(params)["query"]["pages"][0] + page: StrDict = api_get(params)["query"]["pages"][0] if page.get("missing"): raise MissingPage return page -def random_article_list(limit=50): +def random_article_list(limit: int = 50) -> list[StrDict]: + """Get random sample of articles.""" params = { "list": "random", "rnnamespace": "0", "rnlimit": limit, } - return api_get(params)["query"]["random"] + return typing.cast(list[StrDict], api_get(params)["query"]["random"]) def wiki_search(q): @@ -185,23 +193,6 @@ def categorymembers(q: str) -> list[str]: return [i["title"] for i in ret["categorymembers"] if i["title"] != q] -def page_links(titles): # unused - titles = list(titles) - assert titles - params = { - "prop": "links", - "pllimit": 500, - "plnamespace": 0, - "titles": "|".join(titles), - } - ret = api_get(params)["query"] - return dict( - (doc["title"], {l["title"] for l in doc["links"]}) - for doc in ret["pages"].values() - if "links" in doc - ) - - def find_disambig(titles: list[str]) -> list[str]: """Find disambiguation articles in the given list of titles.""" titles = list(titles) @@ -235,7 +226,8 @@ def find_disambig(titles: list[str]) -> list[str]: return disambig -def wiki_redirects(q): # pages that link here +def wiki_redirects(q: str) -> typing.Iterator[str]: + """Pages that link here.""" params = { "list": "backlinks", "blfilterredir": "redirects", @@ -269,7 +261,8 @@ def wiki_backlink(q: str) -> tuple[set[str], set[str]]: return (articles, redirects) -def call_get_diff(title, section_num, section_text): +def call_get_diff(title: str, section_num: int, section_text: str) -> str: + """Get diff from Wikipedia.""" data = { "prop": "revisions", "rvprop": "timestamp", @@ -281,4 +274,4 @@ def call_get_diff(title, section_num, section_text): s = get_session() ret = s.post(get_query_url(), data=data).json() check_for_error(ret) - return ret["query"]["pages"][0]["revisions"][0]["diff"]["body"] + return typing.cast(str, ret["query"]["pages"][0]["revisions"][0]["diff"]["body"])