summaryrefslogtreecommitdiff
path: root/s2-search.py
blob: 1804207b1df95f91db05b457214fc74b0b2deee4 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
import os
import sys
import csv
import subprocess
import time
import random
import re
import json
import click
from s2 import SemanticScholarAPI

'''
s2 search API format:
results
matchedAuthors
matchedPresentations
query
querySuggestions
results
stats
totalPages
totalResults
'''

@click.command()
@click.option('--index', '-n', default=0, help='Index of CSV.')
def fetch_entries(index):
  keys, lines = read_citation_list(index)
  s2 = SemanticScholarAPI()
  for line in lines:
    label = line[0]
    title = re.sub(r'[^-0-9a-zA-Z ]+', '', line[1])
    entry_fn = './datasets/s2/entries/{}.json'.format(title)
    if not os.path.exists(entry_fn):
      results = s2.search(title)
      write_json(dump_fn, results)
      if len(results['results']) == 0:
        print("No results for {}".format(title))
      else:
        print(title)
        write_json(entry_fn, results['results'][0])
      time.sleep(random.randint(10, 20))

def read_citation_list(index=0):
  filename = './datasets/citations.csv'
  if index > 0:
    fn, ext = os.path.splitext(filename)
    filename = fn + '-' + str(index) + ext
  with open(filename, 'r') as f:
    reader = csv.reader(f)
    lines = list(reader)
    keys = lines[0]
    lines = lines[1:]
    return keys, lines

def write_json(fn, data):
  with open(fn, 'w') as outfile:
    json.dump(data, outfile)
def write_csv(fn, keys, rows):
  with open(fn, 'w') as f:
    writer = csv.writer(f)
    if keys is not None:
      writer.writerow(keys)
    for row in rows:
      writer.writerow(row)

if __name__ == '__main__':
  fetch_entries()