# Copyright 2007 Owen Taylor
#
# This file is part of Reinteract and distributed under the terms
# of the BSD license. See the file COPYING in the Reinteract
# distribution for full details.
#
########################################################################

import re

TOKEN_KEYWORD      = 1
TOKEN_NAME         = 2
TOKEN_PUNCTUATION  = 3
TOKEN_COMMENT      = 4
TOKEN_STRING       = 5
TOKEN_CONTINUATION = 6
TOKEN_NUMBER       = 7
TOKEN_JUNK         = 8
TOKEN_LPAREN       = 9
TOKEN_RPAREN       = 10
TOKEN_LSQB         = 11
TOKEN_RSQB         = 11
TOKEN_LBRACE       = 12
TOKEN_RBRACE       = 13
TOKEN_BACKQUOTE    = 14
TOKEN_COLON        = 15
TOKEN_DOT          = 16
TOKEN_EQUAL        = 17
TOKEN_AUGEQUAL     = 18
TOKEN_BUILTIN_CONSTANT = 19

FLAG_OPEN = 1
FLAG_CLOSE = 2

_KEYWORDS = set([ 'and', 'as', 'assert', 'break', 'class', 'continue', 'def',
                  'del', 'elif', 'else', 'except', 'exec', 'finally',  'for',
                  'from', 'global', 'if', 'import', 'in', 'is', 'lambda', 'not',
                  'or', 'pass', 'print', 'raise',  'return', 'try', 'while',
                  'with', 'yield' ])

_IDENTIFIER_TOKENS = {
    'None' : TOKEN_BUILTIN_CONSTANT,
    'True' : TOKEN_BUILTIN_CONSTANT,
    'False' : TOKEN_BUILTIN_CONSTANT
}

_PUNCTUATION_TOKENS = {
    '(' : TOKEN_LPAREN,
    ')' : TOKEN_RPAREN,
    '[' : TOKEN_LSQB,
    ']' : TOKEN_RSQB,
    '{' : TOKEN_LBRACE,
    '}' : TOKEN_RBRACE,
    '`' : TOKEN_BACKQUOTE,
    ':' : TOKEN_COLON,
    '=' : TOKEN_EQUAL,
    '+=' : TOKEN_AUGEQUAL,
    '-=' : TOKEN_AUGEQUAL,
    '*=' : TOKEN_AUGEQUAL,
    '/=' : TOKEN_AUGEQUAL,
    '%=' : TOKEN_AUGEQUAL,
    '&=' : TOKEN_AUGEQUAL,
    '|=' : TOKEN_AUGEQUAL,
    '^=' : TOKEN_AUGEQUAL,
    '<<=' : TOKEN_AUGEQUAL,
    '>>=' : TOKEN_AUGEQUAL,
    '**=' : TOKEN_AUGEQUAL,
    '//=' : TOKEN_AUGEQUAL
}

_PUNCTUATION_MATCH = {
    ')' : '(',
    ']' : '[',
    '}' : '{',
}

_TOKENIZE_RE = re.compile(r"""
# Operators and delimeters
(?P<punctuation>
  [@,:`;~\(\)\[\]\{\}] |
  [+%&|^-]=? |
  \*(?:\*=|\*|=|) |
  /(?:/=|/|=|) |
  <(?:<=|<|=|>|) |
  >(?:>=|>|=|) |
  =(?:=|) |
  !=) |

(?P<comment> \#.*) |                                      # Comment
(?P<identifier> [A-Za-z_][A-Za-z0-9_]*) |                 # Identifier

(?P<string>
  (?:[rR][uU]?|[uU][rR]?|)
  (?P<stringcore>
    (?: '''(?:\\.|[^'\\]|'(?!''))*(?:'''|(?=\\$)|$)) |       # String delimited with '''
    (?: \"""(?:\\.|[^"\\]|"(?!""))*(?:\"""|(?=\\$)|$)) |     # String delimited with \"""
    (?:   '(?:\\.|[^'\\])*(?:'|(?=\\$)|$)) |           # String delimited with '
    (?:   "(?:\\.|[^"\\])*(?:"|(?=\\$)|$))             # String delimited with "
  )
) |

(?P<continuation> \\) |                                   # Line continuation

# A "number-like", possibly invalid expression
(?P<number>
  0[Xx][0-9A-Za-z_]* |

  (?: [0-9] | \.[0-9] )
  [0-9.]*
  (?: [eE][+-]? [0-9A-Za-z_.]* |
      [0-9A-Za-z_.]* )
) |

(?P<dot> \.) |                                            # isolated .
(?P<white> \s+) |                                         # whitespace
(?P<notvalid> [^\s!-\#%->@-~]+) |                         # Not-valid outside of a string... needs to
                                                          # be + not +? to avoid splitting UTF-8
(?P<junk> .+?)                                            # Other junk (vacuum anything notmatched)

""", re.VERBOSE)

_CLOSE_STRING_RE = {
    "'''": re.compile(r"(?:\\.|[^\'\\]|\'(?!\'\'))*(?:(\'\'\')|(?=\\$)|$)"),
    '"""': re.compile(r"(?:\\.|[^\"\\]|\"(?!\"\"))*(?:(\"\"\")|(?=\\$)|$)"),
    "'": re.compile(r"(?:\\.|[^\'\\])*(?:(\')|(?=\\$)|$)"),
    '"': re.compile(r"(?:\\.|[^\"\\])*(?:(\")|(?=\\$)|$)")
 }

# A valid number; the idea is that when tokenizing, we want to keep
# together sequences like 0junk or 0e+a and then mark them as entirely
# invalid rather than breaking them into a "valid" number and a "valid"
# part after that

_NUMBER_RE = re.compile(r"""
^(?:
0j? |                                         # 0 (or complex)
0[Xx][0-9A-Fa-f_]* |                          # Hex 
0[0-7]+ |                                     # Octal
(?:[1-9][0-9]*|0)j? |                         # Decimal (or complex)
(?:(?:[0-9]*\.[0-9]+|[0-9]+\.?)[eE][+-]?[0-9]+ | # Floating point (or complex)
      [0-9]*\.[0-9]+|[0-9]+\.)j?
)$
""", re.VERBOSE)

def tokenize_line(str, stack=None):
    if (stack == None):
        stack = []
    else:
        stack = list(stack)

    tokens = []
    pos = 0
    
    if len(stack) > 0:
        if stack[-1] in _CLOSE_STRING_RE:
            delim = stack[-1]
            match = _CLOSE_STRING_RE[delim].match(str)
            assert(match)

            flags = 0

            if match.group(1):
                flags |= FLAG_CLOSE
                stack.pop()
                
            tokens.append((TOKEN_STRING, match.start(), match.end(), flags))
            
            pos = match.end()
            
    l = len(str)
    while pos < l:
        match = _TOKENIZE_RE.match(str, pos)
        assert(match)
#        print repr(match.group()), match.span(), match.groupdict()
        
        if not match.group('white'):
            flags = 0
            
            token_type = None
            if match.group('punctuation'):
                token_type = TOKEN_PUNCTUATION
                s = match.group()
                if s in _PUNCTUATION_TOKENS:
                    token_type = _PUNCTUATION_TOKENS[s]
                    if token_type == TOKEN_BACKQUOTE:
                        if len(stack) > 0 and stack[-1] == "`":
                            flags |= FLAG_CLOSE
                            stack.pop()
                        else:
                            flags |= FLAG_OPEN
                            stack.append("`")
                    elif s in _PUNCTUATION_MATCH:
                        if len(stack) > 0 and stack[-1] == _PUNCTUATION_MATCH[s]:
                            flags |= FLAG_CLOSE
                            stack.pop()
                        else:
                            token_type = TOKEN_JUNK
                    elif token_type == TOKEN_LPAREN or token_type == TOKEN_LSQB or token_type == TOKEN_LBRACE:
                        flags |= FLAG_OPEN
                        stack.append(s)
            elif match.group('identifier'):
                s = match.group()
                if s in _KEYWORDS:
                    token_type = TOKEN_KEYWORD
                elif s in _IDENTIFIER_TOKENS:
                    token_type = _IDENTIFIER_TOKENS[s]
                else:
                    token_type = TOKEN_NAME
            elif match.group('number'):
                s = match.group()
                if _NUMBER_RE.match(s):
                    token_type = TOKEN_NUMBER
                else:
                    token_type = TOKEN_JUNK
            elif match.group('string'):
                token_type = TOKEN_STRING
                core = match.group('stringcore')
                if core.startswith('"""'):
                    delim = '"""'
                elif core.startswith("'''"):
                    delim = "'''"
                elif core.startswith("'"):
                    delim = "'"
                else:
                    delim = '"'
                if len(core) == len(delim) or \
                   not core.endswith(delim) or \
                   (core[len(core)-len(delim)-1] == '\\' and
                    core[len(core)-len(delim)-2] != '\\'):
                    flags |= FLAG_OPEN
                    stack.append(delim)

            elif match.group('dot'):
                token_type = TOKEN_DOT
            elif match.group('comment'):
                token_type = TOKEN_COMMENT
            elif match.group('continuation'):
                token_type = TOKEN_CONTINUATION
            elif match.group('notvalid') or match.group('junk'):
                token_type = TOKEN_JUNK

            tokens.append((token_type, match.start(), match.end(), flags))
                
        pos = match.end()

    # Catch an unterminated, uncontinued short string, and don't leave it on the stack
    # Would be nice to indicate an error here somehow, but I'm not sure how
    if len(stack) > 0 and (stack[-1] == "'" or stack[-1] == '"') and \
            (len(tokens) == 0  or tokens[-1][0] != TOKEN_CONTINUATION):
        token_type, start, end, flags = tokens[-1]
        flags &= ~FLAG_OPEN
        tokens[-1] = (token_type, start, end, flags)
        stack.pop()
        
    return (tokens, stack)
    
if __name__ == '__main__':
    import sys
    
    failed = False
    
    def expect(str, expected_tokens, in_stack=[], expected_stack=[]):
        tokens, stack = tokenize_line(str, stack=in_stack)
        result = [(token[0], str[token[1]:token[2]]) for token in tokens]
        
        success = True
        if len(tokens) == len(expected_tokens):
            for (t, e) in zip(result, expected_tokens):
                if t != e:
                    success = False
                    break
        else:
            success = False