Zakaria
267 lines
9.2 KiB
Python
267 lines
9.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Search Wikipedia + Wikidata for an entity (person, company, place, concept).
|
|
|
|
Two free APIs:
|
|
- Wikipedia OpenSearch + REST summary endpoint for narrative bio
|
|
- Wikidata SPARQL endpoint for structured facts (birth, employer, awards, etc.)
|
|
|
|
Both are anonymous-access. Useful for resolving who-is-this-entity questions
|
|
and surfacing cross-references that other sources can join against.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import re
|
|
import sys
|
|
import urllib.parse
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).parent))
|
|
from _http import get_json # noqa: E402
|
|
|
|
WP_OPENSEARCH = "https://en.wikipedia.org/w/api.php"
|
|
WP_SUMMARY = "https://en.wikipedia.org/api/rest_v1/page/summary/"
|
|
WD_ACTION = "https://www.wikidata.org/w/api.php"
|
|
|
|
COLUMNS = [
|
|
"source",
|
|
"label",
|
|
"description",
|
|
"qid",
|
|
"wikipedia_title",
|
|
"wikipedia_url",
|
|
"wikidata_url",
|
|
"instance_of",
|
|
"country",
|
|
"occupation",
|
|
"employer",
|
|
"date_of_birth",
|
|
"place_of_birth",
|
|
"summary",
|
|
]
|
|
|
|
|
|
def _wp_search(query: str, limit: int) -> list[dict]:
|
|
params = {
|
|
"action": "opensearch",
|
|
"search": query,
|
|
"limit": str(min(limit, 20)),
|
|
"format": "json",
|
|
}
|
|
url = f"{WP_OPENSEARCH}?{urllib.parse.urlencode(params)}"
|
|
data = get_json(url)
|
|
if not isinstance(data, list) or len(data) < 4:
|
|
return []
|
|
titles, descs, urls = data[1], data[2], data[3]
|
|
out = []
|
|
for i, title in enumerate(titles):
|
|
out.append(
|
|
{
|
|
"title": title,
|
|
"description": descs[i] if i < len(descs) else "",
|
|
"url": urls[i] if i < len(urls) else "",
|
|
}
|
|
)
|
|
return out
|
|
|
|
|
|
def _wp_summary(title: str) -> dict:
|
|
"""Pull the REST summary for a title — short bio, image, type."""
|
|
url = f"{WP_SUMMARY}{urllib.parse.quote(title.replace(' ', '_'))}"
|
|
try:
|
|
return get_json(url) # type: ignore[return-value]
|
|
except Exception as e: # noqa: BLE001
|
|
print(f"Wikipedia summary lookup for {title!r} failed: {e}", file=sys.stderr)
|
|
return {}
|
|
|
|
|
|
def _wd_lookup_by_qid(qid: str) -> dict:
|
|
"""Pull common facts for a QID via Wikidata's Action API (no SPARQL).
|
|
|
|
The Action API is far more lenient on rate-limits than the SPARQL Query
|
|
Service. We get claims as QIDs and then resolve labels in one batch call.
|
|
"""
|
|
# Properties of interest. The Action API returns claims as QIDs or
|
|
# typed literals, so the slot mapping is local-only.
|
|
interesting = {
|
|
"P31": "instance_of",
|
|
"P17": "country", # for orgs / places
|
|
"P27": "country", # for individuals (country of citizenship)
|
|
"P106": "occupation",
|
|
"P108": "employer",
|
|
"P569": "date_of_birth",
|
|
"P19": "place_of_birth",
|
|
}
|
|
params = {
|
|
"action": "wbgetentities",
|
|
"ids": qid,
|
|
"props": "claims",
|
|
"format": "json",
|
|
}
|
|
url = f"{WD_ACTION}?{urllib.parse.urlencode(params)}"
|
|
try:
|
|
data = get_json(url)
|
|
except Exception as e: # noqa: BLE001
|
|
print(f"Wikidata wbgetentities for {qid} failed: {e}", file=sys.stderr)
|
|
return {}
|
|
if not isinstance(data, dict):
|
|
return {}
|
|
claims = (data.get("entities", {}).get(qid, {}) or {}).get("claims", {}) or {}
|
|
|
|
# Collect raw values (QIDs or literals) and remember which slot each
|
|
# came from. Date literals come back as ISO strings; QIDs need a label
|
|
# resolution pass.
|
|
qid_to_slots: dict[str, list[str]] = {}
|
|
facts: dict[str, list[str]] = {}
|
|
for prop_id, slot in interesting.items():
|
|
for claim in claims.get(prop_id, []) or []:
|
|
v = (claim.get("mainsnak", {}) or {}).get("datavalue", {}) or {}
|
|
vtype = v.get("type")
|
|
value = v.get("value")
|
|
if vtype == "wikibase-entityid" and isinstance(value, dict):
|
|
vqid = value.get("id", "")
|
|
if vqid:
|
|
qid_to_slots.setdefault(vqid, [])
|
|
if slot not in qid_to_slots[vqid]:
|
|
qid_to_slots[vqid].append(slot)
|
|
elif vtype == "time" and isinstance(value, dict):
|
|
raw = value.get("time", "") or ""
|
|
# +1955-10-28T00:00:00Z → 1955-10-28
|
|
m = re.search(r"[+-]?(\d{4})-(\d{2})-(\d{2})", raw)
|
|
if m:
|
|
facts.setdefault(slot, []).append(
|
|
f"{m.group(1)}-{m.group(2)}-{m.group(3)}"
|
|
)
|
|
elif vtype == "string":
|
|
facts.setdefault(slot, []).append(str(value))
|
|
|
|
# Resolve labels for all referenced QIDs in one batch (up to 50 at a time).
|
|
qids = list(qid_to_slots)
|
|
for i in range(0, len(qids), 50):
|
|
batch = qids[i : i + 50]
|
|
params = {
|
|
"action": "wbgetentities",
|
|
"ids": "|".join(batch),
|
|
"props": "labels",
|
|
"languages": "en",
|
|
"format": "json",
|
|
}
|
|
url = f"{WD_ACTION}?{urllib.parse.urlencode(params)}"
|
|
try:
|
|
data = get_json(url)
|
|
except Exception as e: # noqa: BLE001
|
|
print(f"Wikidata label batch failed: {e}", file=sys.stderr)
|
|
continue
|
|
if not isinstance(data, dict):
|
|
continue
|
|
ents = data.get("entities", {}) or {}
|
|
for vqid, ent in ents.items():
|
|
label = (ent.get("labels", {}).get("en", {}) or {}).get("value", "") or vqid
|
|
for slot in qid_to_slots.get(vqid, []):
|
|
facts.setdefault(slot, []).append(label)
|
|
|
|
# Deduplicate per slot, preserving order.
|
|
deduped: dict[str, list[str]] = {}
|
|
for slot, vals in facts.items():
|
|
seen = set()
|
|
out = []
|
|
for v in vals:
|
|
if v in seen:
|
|
continue
|
|
seen.add(v)
|
|
out.append(v)
|
|
deduped[slot] = out
|
|
return deduped
|
|
|
|
|
|
def _wd_qid_for_title(title: str) -> str:
|
|
"""Get the Wikidata QID associated with a Wikipedia article title."""
|
|
params = {
|
|
"action": "query",
|
|
"format": "json",
|
|
"prop": "pageprops",
|
|
"ppprop": "wikibase_item",
|
|
"titles": title,
|
|
"redirects": 1,
|
|
}
|
|
url = f"{WP_OPENSEARCH}?{urllib.parse.urlencode(params)}"
|
|
try:
|
|
data = get_json(url)
|
|
except Exception: # noqa: BLE001
|
|
return ""
|
|
if not isinstance(data, dict):
|
|
return ""
|
|
pages = data.get("query", {}).get("pages", {}) or {}
|
|
for page in pages.values():
|
|
qid = (page.get("pageprops") or {}).get("wikibase_item", "")
|
|
if qid:
|
|
return qid
|
|
return ""
|
|
|
|
|
|
def fetch(query: str, limit: int, no_wikidata: bool, out_path: str) -> int:
|
|
hits = _wp_search(query, limit)
|
|
rows: list[dict[str, str]] = []
|
|
for hit in hits[:limit]:
|
|
title = hit.get("title", "")
|
|
if not title:
|
|
continue
|
|
summary = _wp_summary(title)
|
|
qid = _wd_qid_for_title(title) if not no_wikidata else ""
|
|
facts: dict = {}
|
|
if qid:
|
|
facts = _wd_lookup_by_qid(qid)
|
|
rows.append(
|
|
{
|
|
"source": "wikipedia+wikidata" if qid else "wikipedia",
|
|
"label": title,
|
|
"description": (summary.get("description") or hit.get("description") or "").strip(),
|
|
"qid": qid,
|
|
"wikipedia_title": title,
|
|
"wikipedia_url": hit.get("url", ""),
|
|
"wikidata_url": f"https://www.wikidata.org/wiki/{qid}" if qid else "",
|
|
"instance_of": "; ".join(facts.get("instance_of", [])),
|
|
"country": "; ".join(facts.get("country", [])),
|
|
"occupation": "; ".join(facts.get("occupation", [])),
|
|
"employer": "; ".join(facts.get("employer", [])),
|
|
"date_of_birth": "; ".join(facts.get("date_of_birth", []))[:10] if facts.get("date_of_birth") else "",
|
|
"place_of_birth": "; ".join(facts.get("place_of_birth", [])),
|
|
"summary": (summary.get("extract") or "").replace("\n", " ")[:1000],
|
|
}
|
|
)
|
|
|
|
Path(out_path).parent.mkdir(parents=True, exist_ok=True)
|
|
with open(out_path, "w", newline="", encoding="utf-8") as fh:
|
|
w = csv.DictWriter(fh, fieldnames=COLUMNS)
|
|
w.writeheader()
|
|
w.writerows(rows)
|
|
if not rows:
|
|
print(
|
|
f"Wikipedia: 0 articles for query={query!r}. "
|
|
"Private individuals not notable enough for a Wikipedia article "
|
|
"won't appear here (the bar is real).",
|
|
file=sys.stderr,
|
|
)
|
|
return len(rows)
|
|
|
|
|
|
def main() -> int:
|
|
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
p.add_argument("--query", required=True, help="Entity name (person, company, place, concept)")
|
|
p.add_argument("--limit", type=int, default=5)
|
|
p.add_argument(
|
|
"--no-wikidata",
|
|
action="store_true",
|
|
help="Skip the Wikidata SPARQL enrichment (faster, less detail)",
|
|
)
|
|
p.add_argument("--out", required=True)
|
|
a = p.parse_args()
|
|
n = fetch(query=a.query, limit=a.limit, no_wikidata=a.no_wikidata, out_path=a.out)
|
|
print(f"Wrote {n} Wikipedia/Wikidata rows to {a.out}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|