golang_filter

#!/usr/bin/env python2
# -*- coding: utf-8 -*-
from __future__ import print_function

import itertools as it, operator as op, functools as ft
from contextlib import contextmanager
from collections import defaultdict
from tempfile import NamedTemporaryFile
from os.path import dirname, basename, join, exists
from subprocess import Popen, PIPE
import os, sys, re, io, stat


class Actions(object):

	bak_suffix = '.src-bak'


	@classmethod
	def itercode(cls, src):
		'''Iterates over code lines of src (readline func).
			Returns line (verbatim), stripped line, indent (bytes).'''
		for line in iter(src, ''):
			yield line, line.strip(), len(line) - len(line.lstrip('\t'))

	@classmethod
	def check_noncode(cls, ls, src, blocks=True):
		'''Returns if line (stripped) is non-code - comment/space/block/etc.
			Detects string/comment/verbatim blocks only
				with blocks=True, advancing src (readline func).'''
		return not ls\
			or (blocks and cls.skip_block(ls, src))\
			or ls.startswith('//')

	@classmethod
	def skip_block(cls, line, src):
		'Goes over non-code blocks, advancing src (readline func).'
		res = False
		if line.count('`')%2: # string blocks
			for line in iter(src, ''):
				if line.count('`')%2: break
			res = True
		if re.search(r'^\s*/\*\s', line): # comment blocks
			for line in iter(src, ''):
				if re.search(r'^(.*\s+)?\*/\s*$', line): break
			res = True
		if re.search(r'//\+as-is\s*$', line): # verbatim blocks
			for line in iter(src, ''):
				if re.search(r'//-as-is\s*$', line): break
			res = True
		return res


	@classmethod
	def curl(cls, src):
		## Find all places where indentation changes or case-blocks
		cancer, brace_stack = defaultdict(lambda: defaultdict(str)), list()
		pos_prev_tabs, pos_prev = None, src.tell()
		case_block = False
		for line, ls, pos_tabs in cls.itercode(src.readline):
			if cls.check_noncode(ls, src.readline): continue # comments / empty lines
			if cls.skip_block(line, src.readline): continue # block comments/strings
			if pos_prev_tabs is None: # first meaningful line
				pos_prev_tabs, pos_prev = pos_tabs, src.tell() - len(line)
				continue

			if pos_tabs > pos_prev_tabs: # block start (indent++)
				# assert pos_tabs - pos_prev_tabs == 1
				if case_block: case_block = False # "case:" insides
				else:
					pos = src.tell()
					src.seek(pos_prev)
					ls_prev = src.readline().strip()
					src.seek(pos)
					bo, bc = '()' if re.search(r'^(import|var)\s*$', ls_prev) else '{}'
					brace_stack.append((pos_prev, bo, bc))
					cancer[pos_prev]['open'] += bo

			elif pos_tabs < pos_prev_tabs: # block end (indent--)
				for n in xrange(pos_prev_tabs - pos_tabs):
					if not brace_stack: continue # might cause mismatched indent for braces
					pos_open, bo, bc = brace_stack[-1]
					if bo is None and bc is None: # "case:" end
						if not re.search(r'^(case\b.*|default\s*):$', ls):
							brace_stack.pop()
							cancer[pos_prev]['close'] += '}'
						else: case_block = True # next "case:" follows
					else: # block end proper
						brace_stack.pop()
						if re.search(r'^else\b', ls): # special case - "} else {"
							pos = src.tell() - len(line)
							if cancer[pos]['soft_close']:
								assert pos_prev_tabs > pos_tabs\
									and '\t' not in cancer[pos]['soft_close'], cancer[pos]['soft_close']
								cap = cancer[pos]['soft_close']
								cancer[pos]['soft_close'] = '\t' * (pos_prev_tabs - pos_tabs - 1)
								for n in xrange(pos_prev_tabs-1, pos_tabs, -1):
									cancer[pos]['soft_close'] += cap + '\n' + '\t' * (n-1)
							cancer[pos]['soft_close'] += bc
						else: cancer[pos_prev]['close'] += bc

			else:
				if re.search(r'^(case\b.*|default\s*):$', ls): # detect same-indent "case:" blocks
					if not case_block:
						pos = src.tell()
						src.seek(pos_prev)
						ls_prev = src.readline().strip()
						src.seek(pos)
						assert re.search(r'^(switch|select)\b', ls_prev), ls_prev
						brace_stack.append((pos_prev, None, None))
						case_block = True
						cancer[pos_prev]['open'] += '{'
				else: assert not case_block

			pos_prev_tabs, pos_prev = pos_tabs, src.tell() - len(line)

		## Close any remaining blocks
		if brace_stack:
			for pos_open, bo, bc in brace_stack:
				cancer[pos_prev]['close'] += bc

		## Some templating to get rid of ultra-long function sigs
		src.seek(0)
		type_tpls, pos_prev = dict(), src.tell()
		for line in iter(src.readline, ''):
			if line.strip().endswith(','): # handle multiline func signatures
				line = line.strip()
				for line_ext in iter(src.readline, ''):
					line_ext = line_ext.strip()
					line += ' ' + line_ext
					if not line_ext.endswith(','): break
			match = re.search(r'^type (?P<name>\S+) func(?P<sig>\(.+\S)\s*$', line)
			if match: type_tpls[match.group('name')] = match.group('sig')
			match = re.search(r'^func [\w\d_]+<<(?P<name>\S+)>>(?P<ret>\s+\S.*)?\s*$', line)
			if not match:
				match = re.search(r'^\s*\S+\s+:?=\s+func<<(?P<name>\S+)>>(?P<ret>\s+\S.*)?\s*$', line)
			if match:
				if not match.group('ret'): template = type_tpls[match.group('name')]
				else:
					template = type_tpls[match.group('name')]
					template = ( (template.rsplit(')', 2)[0] + ')')\
						if template.endswith(')') else (template.rsplit(')', 1)[0] + ')') )
				cancer[pos_prev]['template'] = '<<{}>>'.format(match.group('name')), template
			pos_prev = src.tell()

		## Bloatify source
		src.seek(0)
		result, pos_prev = list(), src.tell()
		brace_stack, cancer = list(), sorted(cancer.viewitems(), reverse=True)
		tpl_tweaks = list()
		while cancer:
			pos_line, braces = cancer.pop()
			if pos_line is not None:
				# Read to the next brace-line
				assert pos_line >= pos_prev
				result.append(src.read(pos_line - pos_prev))
				assert pos_line == src.tell(), [result[-1][-100:], pos_line, src.tell(), pos_prev]
				line = src.readline()
			else: line = ''
			# Do some templating, if necessary here
			if line and braces.get('template'):
				line_new = line.replace(*braces['template'])
				tpl_tweaks.append((line_new, line))
				line = line_new
			# Append cancer
			bo, bc = it.imap(braces.get, ['open', 'close'])
			pos_tabs, line = len(line) - len(line.lstrip('\t')), line.strip()
			# Keep stack of braces to skip corresponding open/close ones
			if bo:
				for c in [bo, '+', ',', '(']: # long lines, calls and already-curled blocks
					if line.endswith(c):
						bo = None
						break
				brace_stack.append(bo)
			if bc and not brace_stack.pop(): bc = None # opening brace skipped
			if braces.get('soft_close') and brace_stack.pop(): # "} else {"
				line = braces.get('soft_close') + ' ' + line
			# Result line(s) composition
			eol = braces.get('eol', '\n')
			if bc: # special handling for ad-hoc closures, which end in "}()"
				pos = src.tell()
				line_next = src.readline()
				if line_next.strip() != '()': src.seek(pos)
				else:
					pos_next_tabs = len(line_next)
					eol = line_next.lstrip('\t')
					pos_next_tabs -= len(eol) + 1
					pos_line, braces = cancer[-1]
					if pos_line < src.tell():
						if braces.get('close'):
							braces['close'], braces['eol'] = list(braces['close']), ''
							braces['close'][0] = '\t' * pos_next_tabs + braces['close'][0]
						cancer[-1] = None, braces # don't do seek on next iteration
			result.append(
				'\t'*pos_tabs + line
				# + (((' ' if not re.search(r'\[\S*\]\w+\s*$', line) else '') + bo) if bo else '')
				+ ((' ' + bo) if bo else '')
				+ (''.join( ('\n' + '\t'*(pos_tabs-n) + b)
					for n,b in enumerate(bc, 1) ) if bc else '') + eol )
			pos_prev = src.tell()
		result.append(src.read())

		return ''.join(result), tpl_tweaks

	@classmethod
	def curl_postproc(cls, data):
		'Irreversible post-processing'
		return '\n'.join(it.imap(
			ft.partial(re.sub, r'\s*//(\+|-)as-is\s*$', ''),
			''.join(data).splitlines() )) + '\n'


	@classmethod
	def uncurl(cls, src):
		## Find position of every to-be-removed brace
		blanks, pos_prev = defaultdict(bytes), src.tell()
		for line, pos_ls, pos_tabs in cls.itercode(src.readline):
			if cls.check_noncode(pos_ls, src.readline):
				pos_prev = src.tell()
				continue
			for bo,bc in it.izip('{(', '})'):
				if not pos_ls.endswith(bo): continue # opening block-brace at eol
				pos, pos_open_line = src.tell(), pos_prev
				# Find tab or a closing brace
				pos_prev = src.tell()
				for line, ls, line_tabs in cls.itercode(src.readline):
					if not cls.check_noncode(ls, src.readline):
						line_tabs = len(line) - len(line.lstrip('\t'))
						if line_tabs == pos_tabs and re.search(r'^\{}(\s+else\b)?'.format(bc), ls):
							# Found closing brace (and maybe opening one on the same line)
							blanks[pos_open_line] += bo + bc
							blanks[pos_prev] += bo + bc
							break
						elif line_tabs <= pos_tabs: # case in either switch or select, otherwise error
							assert re.search(r'^(switch|select)\b', pos_ls), pos_ls
							assert re.search(r'^(case|default)\b', ls), ls
					pos_prev = src.tell()
				else: raise ValueError('No closing brace found for line: {!r}'.format(line))
				src.seek(pos) # return to next line within the block
			pos_prev = src.tell()

		## Build a clean version of source
		src.seek(0)
		result, pos_prev = list(), src.tell()
		for pos_line, braces in sorted(blanks.viewitems()):
			# Read to the next brace-line
			assert pos_line >= pos_prev
			result.append(src.read(pos_line - pos_prev))
			assert pos_line == src.tell(), [result[-1][-100:], pos_line, src.tell(), pos_prev]
			# Drop braces and spaces near these
			line = src.readline()
			assert any((c in line) for c in braces)
			pos_tabs, line = len(line) - len(line.lstrip('\t')), line.strip().strip(braces).strip()
			if line: result.append('\t'*pos_tabs + line + '\n')
			pos_prev = src.tell()
		result.append(src.read())

		return ''.join(result), list()


	@classmethod
	def replace(cls, path, contents, update_mtime=False, backup=False):
		assert exists(path)
		with NamedTemporaryFile( dir=dirname(path),
				prefix=basename(path)+'.', delete=False ) as tmp:
			pstat = os.stat(path)
			try:
				tmp.write(contents)
				tmp.flush()
				os.fchmod(tmp.fileno(), stat.S_IMODE(pstat.st_mode))
				if not update_mtime:
					os.utime(tmp.name, (pstat.st_mtime, pstat.st_mtime))
				os.fsync(tmp.fileno())
				if backup: os.rename(path, path + cls.bak_suffix)
				os.rename(tmp.name, path)
			finally:
				try: os.unlink(tmp.name)
				except OSError: pass



class NoChange(Exception): pass
class IrreversibleOperation(Exception): pass


def main(argv=None, optz=None, detect_nochange=False):
	if not optz:
		import argparse
		parser = argparse.ArgumentParser(
			description='Pythonify Go sources by removing all brace-cancer.'
				' Operation is perfectly reversible and is checked with "diff -u"'
					' (reverse vs original) on "uncurl" command (unless --force is specified).')
		cmds = parser.add_subparsers(
			title='Supported operations (have their own suboptions as well)')

		@contextmanager
		def subcommand(name, call=None, **kwz):
			cmd = cmds.add_parser(name, **kwz)
			cmd.set_defaults(call=call or name)
			yield cmd

		with subcommand('uncurl', help='Remove curly braces.') as cmd:
			cmd.add_argument('path', help='File to process.')
			cmd.add_argument('-f', '--force', action='store_true',
				help='Dont run "diff" to check whether changes are perfectly reversible.')
			cmd.add_argument('-b', '--backup', action='store_true',
				help=( 'Try to restore "*{}" file instead'
					' of performing operation.'.format(Actions.bak_suffix) ))
		with subcommand('curl', help='Add curly braces.') as cmd:
			cmd.add_argument('path', help='File to process.')
			cmd.add_argument('-f', '--force', action='store_true',
				help='Dont run "diff" to check whether changes are perfectly reversible.')
			cmd.add_argument('-n', '--line', type=int,
				help='Print only specific line from the resulting source.')
			cmd.add_argument('-c', '--context', type=int,
				help='Used with --line will print given number of context lines around it.')
			cmd.add_argument('-p', '--pretty', action='store_true',
				help='With --line, will pretty-print line numbers on the left of each printed line.')
			cmd.add_argument('-b', '--backup', action='store_true',
				help='Create "*{}" file with the original source.'.format(Actions.bak_suffix))
		with subcommand('show', help='Dont change file, just show lines from it.') as cmd:
			cmd.add_argument('path', help='File to process.')
			cmd.add_argument('-n', '--line', type=int,
				help='Print only specific line from the resulting source.')
			cmd.add_argument('-c', '--context', type=int,
				help='Used with --line will print given number of context lines around it.')
			cmd.add_argument('-p', '--pretty', action='store_true',
				help='With --line, will pretty-print line numbers on the left of each printed line.')

		with subcommand('git-smudge', call='git-uncurl',
				help='Same as "uncurl", but git-filter friendly'
					' - works with stdin/stdout, different error handling.') as cmd:
			cmd.add_argument('path', help='Filename. Ignored - git'
				' supplies contents on stdin and expects processing results on stdout.')
			cmd.add_argument('-f', '--force', action='store_true',
				help='Dont run "diff" to check whether changes are perfectly reversible.')
			cmd.add_argument('-b', '--backup', action='store_true',
				help=( 'Try to restore "*{}" file instead'
					' of performing operation.'.format(Actions.bak_suffix) ))
		with subcommand('git-clean', call='git-curl',
				help='Same as "curl", but git-filter friendly'
					' - works with stdin/stdout, different error handling.') as cmd:
			cmd.add_argument('path', help='Filename. Ignored - git'
				' supplies contents on stdin and expects processing results on stdout.')
			cmd.add_argument('-f', '--force', action='store_true',
				help='Dont run "diff" to check whether changes are perfectly reversible.')
			cmd.add_argument('-b', '--backup', action='store_true',
				help='Create "*{}" file with the original source.'.format(Actions.bak_suffix))

		parser.add_argument('-n', '--dry-run', action='store_true',
			help='Dont replace the files, just output mangled source. Ignored with git-* commands.')
		parser.add_argument('-t', '--update-mtime', action='store_true',
			help='Dont restore original timestamps of changed files'
				' (done to work better with make, editors). Ignored with git-* commands.')
		optz = parser.parse_args(sys.argv[1:] if argv is None else argv)

	if optz.call.startswith('git-'):
		git_mode, optz.call = True, optz.call.split('-', 1)[-1]
	else: git_mode = False

	try:
		src, dst = (io.BytesIO(sys.stdin.read()), sys.stdout)\
			if git_mode else ( open(optz.path),
				None if not (optz.dry_run or (optz.call in ['curl', 'show'] and optz.line))
				else sys.stdout )

		if optz.call == 'show': result = src.read()
		else: result, tpl_tweaks = getattr(Actions, optz.call)(src)

		if optz.call in ['curl', 'uncurl'] and not optz.force:
			# Make sure that it's possible to restore original source
			with NamedTemporaryFile() as tmp1,\
					NamedTemporaryFile() as tmp2:
				src.seek(0)
				tmp1.write(src.read())
				tmp1.flush()
				tmp2.write(result)
				tmp2.flush()
				proc = Popen(['diff', '-u', tmp1.name, tmp2.name], stdout=PIPE)
				diff = proc.stdout.read()
				if not diff and not proc.wait():
					if detect_nochange: raise NoChange(result)
				else:
					tmp2.seek(0)
					tmp2.truncate()
					src_rev = io.BytesIO(result)
					result_rev, tpl_tweaks_rev = ( Actions.curl
						if optz.call == 'uncurl' else Actions.uncurl )(src_rev)
					for tweak in tpl_tweaks: # revert templating
						assert tweak[0] in result_rev
						result_rev = result_rev.replace(*tweak)
					tmp2.write(result_rev)
					tmp2.flush()
					proc = Popen(['diff', '-u', tmp1.name, tmp2.name], stdout=PIPE)
					diff = proc.stdout.read()
					if diff or proc.wait():
						raise IrreversibleOperation(diff)

		if optz.call == 'uncurl' and optz.backup: # restore result from backup
			path = optz.path + Actions.bak_suffix
			if exists(path):
				with open(path) as bak: result = bak.read()

		if dst is None:
			Actions.replace( optz.path, result,
				update_mtime=optz.update_mtime,
				backup=(optz.call == 'curl' and optz.backup) )

		else:
			if not git_mode and optz.call in ['curl', 'show'] and optz.line:
				a, b = max(0, ( optz.line - 1
					- (optz.context or 0) )), (optz.line + (optz.context or 0))
				try: result = result.splitlines()[a:b]
				except IndexError:
					print( 'Unable to find specified lines'
						' ({}-{}) in a resulting file'.format(a+1, b+1), file=sys.stderr )
					sys.exit(1)
				result = '\n'.join(
					('{}{}:'.format('> ' if n == optz.line else '  ', n) + line)
					for n, line in enumerate(result, a+1) ) + '\n'

			if optz.call == 'curl':
				if optz.backup: # backup source
					src.seek(0)
					with open(optz.path + Actions.bak_suffix, 'w') as bak: bak.write(src.read())
				if git_mode: # remove processing-related comments from final source
					result = Actions.curl_postproc(result)

			dst.write(result)

	except Exception as err:
		if not git_mode: raise
		# Less verbose errors for git, as there can potentially be a lot of these
		if isinstance(err, IrreversibleOperation): err = 'Irreversible operation'
		print(' -- Failed processing path: {} - {}'.format(optz.path, err), file=sys.stderr)
		sys.exit(1)


if __name__ == '__main__':
	try: main()
	except IrreversibleOperation as err:
		print('Failed to do reverse-transformation flawlessly', file=sys.stderr, end='\n\n')
		print(err.args[0], file=sys.stderr, end='')
		sys.exit(1)