bin/edit-csv-as-yaml

   1 #!/usr/bin/env python
   2 """Quick, dirty script that will convert a csv file to yaml, spawn an editor
   3 for you to fiddle with it, then convert back to csv and replace the original
   4 file.
   5
   6 Run me as: $0 some_file.csv
   7
   8 The editor used is $EDITOR, of course.
   9
  10 This script is not guaranteed to be even remotely reliable, so consider only
  11 using it on files in source control.
  12 """
  13
  14 import codecs
  15 import csv
  16 import os
  17 import subprocess
  18 import sys
  19 import tempfile
  20
  21 try:
  22     import yaml
  23 except ImportError:
  24     sys.stderr.write("Please install PyYAML.\n")
  25     sys.exit(13)
  26
  27 infilename, = sys.argv[1:]
  28
  29 data = []
  30 with open(infilename) as infile:
  31     reader = csv.reader(infile, lineterminator='\n')
  32     column_names = [unicode(column) for column in next(reader)]
  33
  34     # Read data...
  35     for row in reader:
  36         datum = dict()
  37         for col, value in zip(column_names, row):
  38             datum[col] = value.decode('utf-8')
  39
  40         data.append(datum)
  41
  42
  43 # Monkeypatch yaml to use > syntax for multiline text; easier to edit
  44 from yaml.emitter import Emitter
  45 orig_choose_scalar_style = Emitter.choose_scalar_style
  46 def new_choose_scalar_style(self):
  47     if self.analysis is None:
  48         self.analysis = self.analyze_scalar(self.event.value)
  49     if self.analysis.multiline:
  50         return '>'
  51     return orig_choose_scalar_style(self)
  52 Emitter.choose_scalar_style = new_choose_scalar_style
  53
  54 # Write to a tempfile
  55 with tempfile.NamedTemporaryFile(suffix='.yml') as tmp:
  56     yaml.safe_dump(data, tmp,
  57         default_flow_style=False,
  58         allow_unicode=True,
  59         indent=4,
  60     )
  61     del data  # reclaim rams!
  62
  63     error_line = ''  # used on errors
  64     while True:
  65         args = [os.environ['EDITOR'], tmp.name]
  66         if 'vim' in os.environ['EDITOR']:
  67             # vim has an arg for jumping to a line:
  68             args.append("+{0}".format(error_line))
  69
  70         # Run the user's editor and wait for it to close
  71         subprocess.Popen(args).wait()
  72         tmp.seek(0)
  73
  74         try:
  75             new_data = yaml.safe_load(tmp)
  76             break
  77         except yaml.YAMLError as e:
  78             if hasattr(e, 'problem_mark'):
  79                 error_line = e.problem_mark.line + 1
  80             else:
  81                 error_line = ''
  82
  83             print
  84             print "Oh my god what have you done:"
  85             print
  86             print str(e)
  87             print
  88             print "Press Enter to try again, or I guess ctrl-c to bail."
  89             raw_input()
  90
  91 with open(infilename, 'wb') as outfile:
  92     writer = csv.writer(outfile, lineterminator='\n')
  93     writer.writerow([ column.encode('utf8') for column in column_names ])
  94
  95     for datum in new_data:
  96         writer.writerow([
  97             datum[column].encode('utf8') for column in column_names
  98         ])