test_get_complete_pos.py

import ast
import logging
import re

# import test_get_complete_position_original as orig

logger = logging.getLogger('test_get_complete_position')
logger.propagate = False

imports = re.compile(
    r'^\s*(from\s+\.*\w+(\.\w+)*\s+import\s+(\w+,\s+)*|import\s+)')
split_pattern = re.compile(r'''[^= \r\n*()@-]''')
keyword = re.compile('[A-Za-z0-9_]')
keyword_or_brackets = re.compile('[A-Za-z0-9_[\]]')
opening = re.compile('^(.*\[)[A-Za-z_''".]')
splitchars = frozenset('= \r\n*()@-')


def parses(code):
    try:
        ast.parse(code)
        return code
    except SyntaxError:
        pass

    while opening.match(code):
        prefix, _, code = code.partition('[')
        try:
            ast.parse(prefix)
            return prefix + '['
        except SyntaxError:
            continue
    return ''


class Source:

    def get_pos(self, context):
        line = context['line']
        col = context['col']
        start = col

        # Check if the cursor is in an incomplete string
        for char in ("'", '"'):
            bracket = line[:start].rfind('[' + char)
            if bracket > 0:
                # Make sure the quote next to the bracket is the last
                # quote before the cursor
                if bracket + 1 == line[:start].rfind(char):
                    logger.debug('starts with %r', char)
                    start = bracket
                    break

        stack = []
        while start > 0:
            char = line[start - 1]
            logger.debug('start=%r char=%r stack=%r', start, char, stack)
            if stack and stack[-1] in ("'", '"') and char == stack[-1]:
                stack.pop()
            elif char in '''"'])}''':
                stack.append(char)
            elif char in '[({':
                if not stack or stack[-1] != '])}'['[({'.index(char)]:
                    break
                stack.pop()
            elif not stack and char in splitchars:
                break
            start -= 1

        return -1 if start == col else start

    def get_complete_position(self, context):
        line = context['line']
        col = context['col']
        logger.debug('line: %r', line[:col])

        # return immediately for imports
        if imports.match(context['input']):
            logger.debug('matches imports')
            start = len(context['input'])
            while start > 0 and re.match('[._A-Za-z0-9]',
                                         context['input'][start - 1]):
                start -= 1
            return start

        # locate the start of the word
        start = col - 1
        if start == 0 or (len(line) == start and
                          not split_pattern.match(line[start - 2]) and
                          not (start >= 2 and
                               keyword.match(line[start - 3])) and
                          line[start - 3:start - 1] != '].'):
            start = -1
            return start

        start = len(line[:col]) - 1
        bracket_level = 0
        char = line[start - 1]
        char2 = line[start - 2] if start >= 2 else ''
        while start > 0 and (
                split_pattern.match(char) or
                (char == '.' and keyword.match(char2)) or
                (char in '-[' and char2 == '[') or
                (char and line[start - 2:start] == '].')):
            if char == '[':
                if start == 1 or not keyword_or_brackets.match(char2):
                    break
                bracket_level += 1
            elif char == ']':
                bracket_level -= 1
            start -= 1
            char = line[start - 1]
            char2 = line[start - 2] if start >= 2 else ''

        logger.debug('bracket level: %d', bracket_level)
        while bracket_level > 0 and opening.match(line[start:col]):
            prefix = parses(line[start:col])
            if prefix:
                logger.debug('removing %r at %d', prefix, start)
                start += len(prefix)
                bracket_level -= 1
            else:
                break

        base = line[start:col]
        if not bracket_level and base.endswith(' '):
            if not re.match(r'^\s*(?:from|import).*\s+$', line):
                return -1

        logger.debug('bracket level: %d', bracket_level)
        return start


source = Source()


def test(line):
    line, ans = line.rsplit(maxsplit=1)
    ans = int(ans)

    context = dict(
        line=line,
        col=len(line),
        input=line,
    )
    logger.info('completing %r', line)
    start = source.get_complete_position(context)
    start2 = source.get_pos(context)
    assert start2 == ans, (line, ans, start2)
    logger.debug('%s, %s', start, start2)
    logger.info('start=%r end=%r', line[:start2], line[start2:])
    # assert start == orig.source.get_complete_position(context)
    return start


def main():

    cases = '''
data[test[0]]['test 0
items[items['enum'] < 50].x 0
items[items['test foo 6
data3['aspernatur'][' 0
foo['test()']. 0
test[0]. 0

this = test[0]['foob 7
(items[x + 1]. 1
test[ -1
foo = a + 'test' + bar 19
{test.foo 1
'test'.split 0
('test'.split 1
('test'.split('x 14
('test'.split('x' + [(te 22
line(data2.t[:, 0], ax=ax 23
    @test. 5
'''.strip().splitlines()

    for case in filter(None, cases):
        test(case)


if __name__ == '__main__':
    main()