Adding ES6 {a}=1 support for python using tokenizer

README.md

It would be cool if there's webpack for python, since sometimes I need a custom syntax for the python.

exec transform('{a, b} = {"a": 1, "b": 1}')
# a = 1
# b = 1

transform.py

"""
ES6 {} assign for python (hacky, since I only use a lexer for python which is builtin)
{a, b} = {'a': 1, 'b': 1}
->
__tmp_variable1 = {'a': 1, 'b': 1}
a = __tmp_variable1['a']
b = __tmp_variable1['b']
del __tmp_variable1

1. Detect {a,b}=(variable) pattern
(tokenizer works?
  - how to detect WHITESPACES (since it's crucial)
  -> no need to detect since I can use ;? untokenize works)
2. Transform to the pattern above
3. Test
"""

import tokenize
import StringIO
import token as tokenDict

# token map

tokenMap = {y: x for x, y in tokenDict.__dict__.items() if type(y) == int}

tokenMap[tokenize.COMMENT] = 'COMMENT'
tokenMap[tokenize.NL] = 'NL'

# string to readline()


def line_reader(x):
    io = StringIO.StringIO(x)

    def f():
        s = io.readline()
        return s
    return f


def isNewlineOrIndent(data):
    type, token = data[:2]
    result = type in (
        tokenize.NL, tokenDict.INDENT) or token in startOfStatementChar
    return result


class ES6State(object):
    def __init__(self, index, firstToken):
        self.started = True
        self.dstTokens = [firstToken]
        self.valueTokens = []
        self.index = index
        self.dictionaryEnded = False
        self.state = 'start'
        self.completed = False

    def allowed(self, data):
        # { name [, name] } = lol
        # TODO: support alternative-name data like {a: b} = {a: 1} -> b = 1
        type, token, spos, epos, line = data
        prevState = self.state

        # If it's really a es6 assign statement and it has newline, it'll be transformed.
        if self.completed == True:
            if not isNewlineOrIndent(data):
                self.valueTokens.append(data)
            return True

        else:
            self.dstTokens.append(data)

            # {a <-, b <-}
            if type in (tokenDict.NAME, ):
                self.state = 'name'
                return prevState in {'start', 'comma'}

            # {,}=
            if token == ',':
                self.state = 'comma'
                return prevState == 'name'
            elif token == '}':
                if self.dictionaryEnded:
                    return False
                self.dictionaryEnded = True
                return True
            elif token == '=':
                self.completed = True
                return True

            # otherwise return False
            return False

    def getTransformedTokens(self, tokens):
        varName = '__temp_var_es6assign__'
        firstToken = self.dstTokens[0]
        expr = '%s=%s' % (varName, ''.join(x[1] for x in self.valueTokens))
        for name in (x for x in self.dstTokens if x[0] == tokenDict.NAME):
            expr += ';%s=%s[%r]' % (name[1], varName, name[1])
        expr += ';del %s' % varName
        assignToken = (tokenDict.OP, expr, firstToken[2], firstToken[3], None)
        tokens.append(assignToken)
        pass

    pass

# source to parse


startOfStatementChar = '\n;\t'


def transform(source):
    startOfStatement = True
    es6state = None
    tokens = []  # storage to save tokens
    es6AssignTokens = None

    for data in tokenize.generate_tokens(line_reader(source)):
        type, token, spos, epos, line = data
        index = len(tokens)
        tokenOverriden = False

        # print tokenMap[type], `token`

        if es6state:
            if not es6state.allowed(data):
                es6state = None

        if startOfStatement and token == '{':
            es6state = ES6State(index, data)

        if isNewlineOrIndent(data):
            if es6state and es6state.completed:  # processing es6 assign
                tokens = tokens[:es6state.index]
                # transformed operators
                es6state.getTransformedTokens(tokens)

            startOfStatement = True
            es6state = None
        else:
            startOfStatement = False

        if not tokenOverriden:
            tokens.append(data)

    return tokenize.untokenize(tokens)


if __name__ == '__main__':
    originalCode = r'''
# Case 1: it's a dictionary/set literal, not es6 assign
{'lol', 'lol!!'}
b = 1
{'a': 1, 'b': [1]};{'a': 1, b: 1}

# Case 2: it's a es6 assign statement
{my_name, lol} = {'my_name': 1, 'lol': 2}
if True:
	{b} = {'b': 100}
{b}=1;
{c}=1 # comment!
	'''

    transformedCode = transform(originalCode)

    LINE = '=' * 80
    print 'Original code:'
    print LINE
    print originalCode
    print LINE
    print 'Transformed code:'
    print LINE
    print transformedCode
    print LINE

    class TracerDict(dict):
        def __init__(self):
            self.items = {}

        def __setitem__(self, x, y):
            print 'assigning', x, 'to', `y`
            self.items[x] = y

        def __getitem__(self, x):
            return self.items[x]

        def __delitem__(self, x):
            del self.items[x]

    print 'Executing code...'
    print LINE
    try:
        exec(transformedCode, TracerDict())
    except TypeError as e:
        assert str(e) == "'int' object has no attribute '__getitem__'"
        print e.__class__.__name__, e
        print 'Success!'

    exec transform('{a, b} = {"a": 1, "b": 1}')
    print a, b