gistfile1.py

#!/usr/bin/python
import os, sys
import subprocess
import readline
from collections import defaultdict

TABSIZE = 2

class TaskDefDict(defaultdict):
  def __missing__(self, key):
    self[key] = [(-TABSIZE, 'def %s():' % key)]
    self[key].append((0, "require('hosts', provided_by=[vagrant])"))
    return self[key]

class FabWrapper(object):
  def __init__(self, args, fname):
    self.subproc = subprocess.Popen(sys.argv[1:], stdin=subprocess.PIPE)
    self.outfile = open(fname, 'w')
    self.write_initial_env()
    self.taskdefs = TaskDefDict()
    self.task_indents = defaultdict(int)
    self.current_task = 'setup_vagrant'
    self.recording = True

  def __enter__(self):
    return self
  def __exit__(self, *args):
    for taskdef in self.taskdefs.itervalues():
      for indent, line in taskdef:
        self.outfile.write((' ' * (indent + TABSIZE)) + line + '\n')
      self.outfile.write('\n')

  def write_initial_env(self):
    self.outfile.write("""
def vagrant():
  raw_ssh_config = subprocess.Popen(['vagrant', 'ssh-config'],
                                    stdout=subprocess.PIPE).communicate()[0]
  ssh_config = dict([l.strip().split() for l in raw_ssh_config.split('\n') if l])
  env.user = ssh_config['User']
  env.hosts = ['%s:%s' % (ssh_config['HostName'], ssh_config['Port'])]
  env.key_filename = ssh_config['IdentityFile']

""")

  def process_directive(self, directive):
    parts = directive.split()
    if parts[0] == 'task':
      self.current_task = parts[1]
    elif parts[0] == 'off':
      self.recording = False
    elif parts[1] == 'on':
      self.recording = True
    elif parts[2] == 'uncd':
      self.decrease_indent()

  def append_line(self, line):
    self.taskdefs[self.current_task].append(
      (self.task_indents[self.current_task], line))

  def modify_indent(self, amt):
    self.task_indents[self.current_task] += amt
  increase_indent = lambda self: self.modify_indent(TABSIZE)
  decrease_indent = lambda self: self.modify_indent(-TABSIZE)

  def process_command(self, cmd):
    if cmd.startswith('cd'):
      self.append_line("with cd('%s'):" % cmd[len('cd '):])
      self.increase_indent
    elif cmd.startswith('sudo'):
      self.append_line("sudo('%s')" % cmd[len('sudo '):])
    else:
      self.append_line("run('%s')" % cmd)

  def readloop(self):
    while True:
      try:
        line = raw_input()
        if line.startswith('#'):
          self.process_directive(line[1:])
        elif self.recording:
          self.process_command(line)
        self.subproc.stdin.write(line + '\n')
      except EOFError:
        break
      except KeyboardInterrupt:
        break

with FabWrapper(sys.argv[1:], 'fabfile.py') as wrapper:
  wrapper.readloop()