#!/usr/bin/python

# Strip trailing whitespace ONLY in modified sections in the latest commit.
# Usage:
# <edit files>
# $ git commit
# $ git un-diff-whitespace
# Copyright 2010 Colin Walters <walters@verbum.org>
# Licensed under the new-BSD license (http://www.opensource.org/licenses/bsd-license.php)

import os,sys,re,subprocess,hashlib

hunk_re = re.compile(r'^@@ -\d+,\d+ \+(\d+),(\d+) @@')
ws_re = re.compile(r'^(.*?)[ \t]+$')
file_re = re.compile(r'^\+\+\+ b/(.+)')

def hashfile(path):
  h = hashlib.sha1()
  f = open(path)
  b = f.read(8192)
  while b != '':
    h.update(b)
    b = f.read(8192)
  f.close()
  return h.hexdigest()

def strip_whitespace_regions(filename, hunks):
  if len(hunks) == 0:
    return False
  filename_t = filename + '.tmp'
  tmpf = open(filename_t, 'w')
  f = open(filename)
  i = 0
  hunk = hunks[0]
  for line in f:
    if (hunk is None) or (i < hunk[0]):
      tmpf.write(line)
    elif i < (hunk[0] + hunk[1]):
      nline = ws_re.sub(r'\1', line) 
      tmpf.write(nline) 
    else:
      hunks = hunks[1:]
      if len(hunks) == 0:
        hunk = None
      else:
        hunk = hunks[0] 
      tmpf.write(line)
    i = i+1
  tmpf.close()
  f.close()
  old = hashfile(filename)
  new = hashfile(filename_t)
  if old != new:
    os.rename(filename_t, filename)
    print "wrote " + filename
    return True
  else:
    os.unlink(filename_t)
    return False

def main():
  toplevel = subprocess.Popen(['git', 'rev-parse', '--show-toplevel'], stdout=subprocess.PIPE).communicate()[0]
  os.chdir(toplevel[:-1])
  diff = subprocess.Popen(['git', 'diff'], stdout=subprocess.PIPE).communicate()[0]
  if diff != '':
    print "You have unstaged changes; Commit them or 'git stash'"
    sys.exit(1)
  commit_proc = subprocess.Popen(['git', 'show', '-p', 'HEAD'], stdout=subprocess.PIPE,
                                 stderr=sys.stderr)
  diff_lines = commit_proc.stdout.readlines()
  commit_proc.wait()
  curfile = None
  hunks = []
  files = []
  modified_files = []
  for line in diff_lines:
    if curfile is not None:
      match = hunk_re.match(line) 
      if match:
        hunks.append((int(match.group(1)),int(match.group(2))))
        continue
    match = file_re.match(line)
    if match:
      filename = match.group(1)
      modified = strip_whitespace_regions(curfile, hunks)
      curfile = filename
      files.append(filename)
      if modified:
        modified_files.append(filename)
      hunks = []
  if curfile:
    modified = strip_whitespace_regions(curfile, hunks)
    if modified:
      modified_files.append(curfile)
  if modified_files:
    add_args = ['git', 'add']
    add_args.extend(modified_files)
    subprocess.check_call(add_args, stdout=sys.stdout, stderr=sys.stderr)
  else:
    print "No trailing whitespace found in HEAD"

if __name__ == '__main__':
  main()