summaryrefslogtreecommitdiff
path: root/util.py
blob: 9f3214658f8a72faa6ff21238ba886ce3e72afb8 (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
import os
import csv
import simplejson as json

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 read_csv(fn, keys=True):
  with open(fn, 'r') as f:
    reader = csv.reader(f)
    lines = list(reader)
    if keys:
      keys = lines[0]
      lines = lines[1:]
      return keys, lines
    return lines

def read_json(fn):
  with open(fn, 'r') as json_file:
    return json.load(json_file)

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)