]> code.communitydata.science - covid19.git/blob - keywords/src/wikidata_translations.py
1ac956c749f0a774a3605faa2ac92db2d33a5ff2
[covid19.git] / keywords / src / wikidata_translations.py
1 from wikidata_api_calls import run_sparql_query
2 from itertools import chain, islice
3 import csv
4 from json import JSONDecodeError
5 from os import path
6
7 class LabelData:
8     __slots__ = ['entityid','label','langcode','is_alt']
9
10     def __init__(self, wd_res, is_alt):
11         obj = wd_res.get('label',None)
12         self.label = obj.get('value',None)
13         self.langcode = obj.get('xml:lang',None)
14         self.entityid = wd_res.get('entity',None).get('value',None)
15         self.is_alt = is_alt
16
17     def to_list(self):
18         return [self.entityid,
19                 self.label,
20                 self.langcode,
21                 self.is_alt]
22
23 def GetAllLabels(in_csvs, outfile, topNs):
24
25     def load_entity_ids(in_csv, topN=5):
26         with open(in_csv,'r',newline='') as infile:
27             reader = list(csv.DictReader(infile))
28             for row in reader:
29                 if int(row['search_position']) < topN:
30                     yield row["entityid"]
31
32     ids = set(chain(* map(lambda in_csv, topN: load_entity_ids(in_csv, topN), in_csvs, topNs)))
33
34     labeldata = GetEntityLabels(ids)
35
36     with open(outfile, 'w', newline='') as of:
37         writer = csv.writer(of)
38         writer.writerow(LabelData.__slots__)
39         writer.writerows(map(LabelData.to_list,labeldata))
40
41     
42 def GetEntityLabels(entityids):
43
44     def run_query_and_parse(query, is_alt):
45         results = run_sparql_query(query)
46         try:
47             jobj = results.json()
48
49             res = jobj.get('results',None)
50             if res is not None:
51                 res = res.get('bindings',None)
52             if res is None:
53                 raise requests.APIError(f"got invalid response from wikidata for {query % entityid}")
54
55             for info in res:
56                 yield LabelData(info, is_alt)
57
58         except JSONDecodeError as e:
59             print(e)
60             print(query)
61             
62     def prep_query(query, prop, entityids):
63         values = ' '.join(('wd:{0}'.format(id) for id in entityids))
64         return query.format(prop, values)
65     
66     base_query = """
67     SELECT DISTINCT ?entity ?label WHERE {{
68     ?entity {0} ?label;
69     VALUES ?entity  {{ {1} }}
70     }}"""
71
72     # we can't get all the entities at once. how about 100 at a time?
73     chunksize = 100
74     entityids = (id for id in entityids)
75     chunk = list(islice(entityids, chunksize))
76     calls = []
77     while len(chunk) > 0:
78         label_query = prep_query(base_query, "rdfs:label", chunk)
79         altLabel_query = prep_query(base_query, "skos:altLabel", chunk)
80         label_results = run_query_and_parse(label_query,  is_alt=False)
81         altLabel_results = run_query_and_parse(altLabel_query, is_alt=True)
82         calls.extend([label_results, altLabel_results])
83         chunk = list(islice(entityids, chunksize))
84
85     return chain(*calls)
86         
87
88 def find_new_output_file(output, i = 1):
89     if path.exists(output):
90         name, ext = path.splitext(output)
91
92         return find_new_output_file(f"{name}_{i}.{ext}", i+1)
93     else:
94         return output
95
96 if __name__ == "__main__":
97     import argparse
98     parser = argparse.ArgumentParser("Use wikidata to find transliterations of terms")
99     parser.add_argument('inputs', type=str, nargs='+', help='one or more files to read. the inputs are generated by wikidata_search.py')
100     parser.add_argument('--topN', type=int, nargs='+', help='limit number of wikidata search results to use, can pass one arg for each source.')
101     parser.add_argument('--output', type=str, help='an output file. defaults to stdout',default=20)
102
103     args = parser.parse_args()
104
105     output = find_new_output_file(args.output)
106
107     GetAllLabels(args.inputs, output, topNs=args.topN)

Community Data Science Collective || Want to submit a patch?