sort-ontologies.py

#!/usr/bin/env python3

import csv
import sys
import yaml

from argparse import ArgumentParser


def main(args):
  parser = ArgumentParser(description='''
  Takes a YAML file containing information for various ontologies and a metadata file specifying
  the sorting order for ontologies, and then produces a sorted version input YAML''')
  parser.add_argument('unsorted_yaml', type=str,
                      help='Unsorted YAML file containing information for ontologies')
  parser.add_argument('metadata_grid', type=str,
                      help='CSV or TSV file containing metadata information for ontologies')
  parser.add_argument('output_yaml', type=str,
                      help='Name of output YAML file that will contain sorted ontology information')
  args = parser.parse_args()

  data_file = args.unsorted_yaml
  grid = args.metadata_grid
  output = args.output_yaml

  sort_order = get_sort_order(grid)
  data = load_data(data_file)
  data = sort_ontologies(data, sort_order)
  write_data(data, output)


def get_sort_order(grid):
  '''Given the path to the metadata grid (CSV or TSV), extract the order of
  ontologies from the grid. Return the list of ontology IDs in that order.'''
  sort_order = []
  if '.csv' in grid:
    separator = ','
  elif '.tsv' or '.txt' in grid:
    separator = '\t'
  else:
    print('%s must be tab- or comma-separated.', file=sys.stderr)
    sys.exit(1)
  with open(grid, 'r') as f:
    reader = csv.reader(f, delimiter=separator)
    # Ignore the header row:
    next(reader)
    for row in reader:
      # Ontology IDs are in the first column of the CSV/TSV. We simply pull them out of each line
      # in the file. Their ordering in the file is the sort ordering we are looking for:
      sort_order.append(row[0])
  return sort_order


def load_data(data_file):
  '''Given a YAML file, load the data into a dictionary.'''
  stream = open(data_file, 'r')
  data = yaml.load(stream, Loader=yaml.SafeLoader)
  return data


def sort_ontologies(data, sort_order):
  '''Given the ontologies data as a dictionary and the list of ontologies in
  proper sort order, return the sorted data.'''
  ontologies = []
  for ont_id in sort_order:
    # We assume that ontology ids are unique:
    ont = [ont for ont in data['ontologies'] if ont['id'] == ont_id].pop()
    ontologies.append(ont)
  data['ontologies'] = ontologies
  return data


def write_data(data, output):
  '''Given the ontologies data as a dictionary and an output YAML file to
  write to, write the data to the file. '''
  yaml_str = yaml.dump(data)
  with open(output, 'w') as f:
    f.write(yaml_str)


if __name__ == '__main__':
  main(sys.argv)