<s> from __future__ import print_function <EOL> import sys <EOL> import sqlparse <EOL> from sqlparse . sql import Comparison , Identifier , Where <EOL> from . parseutils import last_word , extract_tables , find_prev_keyword <EOL> from . special import parse_special_command <EOL> PY2 = sys . version_info [ <NUM_LIT:0> ] == <NUM_LIT:2> <EOL> PY3 = sys . version_info [ <NUM_LIT:0> ] == <NUM_LIT:3> <EOL> if PY3 : <EOL> string_types = str <EOL> else : <EOL> string_types = basestring <EOL> def suggest_type ( full_text , text_before_cursor ) : <EOL> """<STR_LIT>""" <EOL> word_before_cursor = last_word ( text_before_cursor , <EOL> include = '<STR_LIT>' ) <EOL> identifier = None <EOL> if word_before_cursor : <EOL> if word_before_cursor [ - <NUM_LIT:1> ] == '<STR_LIT:(>' or word_before_cursor [ <NUM_LIT:0> ] == '<STR_LIT:\\>' : <EOL> parsed = sqlparse . parse ( text_before_cursor ) <EOL> else : <EOL> parsed = sqlparse . parse ( <EOL> text_before_cursor [ : - len ( word_before_cursor ) ] ) <EOL> p = sqlparse . parse ( word_before_cursor ) [ <NUM_LIT:0> ] <EOL> if p . tokens and isinstance ( p . tokens [ <NUM_LIT:0> ] , Identifier ) : <EOL> identifier = p . tokens [ <NUM_LIT:0> ] <EOL> else : <EOL> parsed = sqlparse . parse ( text_before_cursor ) <EOL> if len ( parsed ) > <NUM_LIT:1> : <EOL> current_pos = len ( text_before_cursor ) <EOL> stmt_start , stmt_end = <NUM_LIT:0> , <NUM_LIT:0> <EOL> for statement in parsed : <EOL> stmt_len = len ( statement . to_unicode ( ) ) <EOL> stmt_start , stmt_end = stmt_end , stmt_end + stmt_len <EOL> if stmt_end >= current_pos : <EOL> text_before_cursor = full_text [ stmt_start : current_pos ] <EOL> full_text = full_text [ stmt_start : ] <EOL> break <EOL> elif parsed : <EOL> statement = parsed [ <NUM_LIT:0> ] <EOL> else : <EOL> statement = None <EOL> if statement : <EOL> tok1 = statement . token_first ( ) <EOL> if tok1 and tok1 . value == '<STR_LIT:\\>' : <EOL> return suggest_special ( text_before_cursor ) <EOL> last_token = statement and statement . token_prev ( len ( statement . tokens ) ) or '<STR_LIT>' <EOL> return suggest_based_on_last_token ( last_token , text_before_cursor , <EOL> full_text , identifier ) <EOL> def suggest_special ( text ) : <EOL> text = text . lstrip ( ) <EOL> cmd , arg = parse_special_command ( text ) <EOL> if cmd == text : <EOL> return [ { '<STR_LIT:type>' : '<STR_LIT>' } ] <EOL> if cmd in ( '<STR_LIT>' , '<STR_LIT>' ) : <EOL> return [ { '<STR_LIT:type>' : '<STR_LIT>' } ] <EOL> if cmd in ( '<STR_LIT>' ) : <EOL> return [ { '<STR_LIT:type>' : '<STR_LIT>' } ] <EOL> if cmd in [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] : <EOL> return [ { '<STR_LIT:type>' : '<STR_LIT>' } ] <EOL> if cmd in [ '<STR_LIT>' ] : <EOL> return [ <EOL> { '<STR_LIT:type>' : '<STR_LIT>' , '<STR_LIT>' : [ ] } , <EOL> { '<STR_LIT:type>' : '<STR_LIT>' , '<STR_LIT>' : [ ] } , <EOL> { '<STR_LIT:type>' : '<STR_LIT>' } , <EOL> ] <EOL> return [ { '<STR_LIT:type>' : '<STR_LIT>' } , { '<STR_LIT:type>' : '<STR_LIT>' } ] <EOL> def suggest_based_on_last_token ( token , text_before_cursor , full_text , identifier ) : <EOL> if isinstance ( token , string_types ) : <EOL> token_v = token . lower ( ) <EOL> elif isinstance ( token , Comparison ) : <EOL> token_v = token . tokens [ - <NUM_LIT:1> ] . value . lower ( ) <EOL> elif isinstance ( token , Where ) : <EOL> prev_keyword , text_before_cursor = find_prev_keyword ( text_before_cursor ) <EOL> return suggest_based_on_last_token ( prev_keyword , text_before_cursor , <EOL> full_text , identifier ) <EOL> else : <EOL> token_v = token . value . lower ( ) <EOL> if not token : <EOL> return [ { '<STR_LIT:type>' : '<STR_LIT>' } , { '<STR_LIT:type>' : '<STR_LIT>' } ] <EOL> elif token_v . endswith ( '<STR_LIT:(>' ) : <EOL> p = sqlparse . parse ( text_before_cursor ) [ <NUM_LIT:0> ] <EOL> if p . tokens and isinstance ( p . tokens [ - <NUM_LIT:1> ] , Where ) : <EOL> column_suggestions = suggest_based_on_last_token ( '<STR_LIT>' , <EOL> text_before_cursor , full_text , identifier ) <EOL> where = p . tokens [ - <NUM_LIT:1> ] <EOL> prev_tok = where . token_prev ( len ( where . tokens ) - <NUM_LIT:1> ) <EOL> if isinstance ( prev_tok , Comparison ) : <EOL> prev_tok = prev_tok . tokens [ - <NUM_LIT:1> ] <EOL> prev_tok = prev_tok . value . lower ( ) <EOL> if prev_tok == '<STR_LIT>' : <EOL> return [ { '<STR_LIT:type>' : '<STR_LIT>' } ] <EOL> else : <EOL> return column_suggestions <EOL> prev_tok = p . token_prev ( len ( p . tokens ) - <NUM_LIT:1> ) <EOL> if prev_tok and prev_tok . value and prev_tok . value . lower ( ) == '<STR_LIT>' : <EOL> tables = extract_tables ( full_text ) <EOL> return [ { '<STR_LIT:type>' : '<STR_LIT>' , '<STR_LIT>' : tables , '<STR_LIT>' : True } ] <EOL> elif p . token_first ( ) . value . lower ( ) == '<STR_LIT>' : <EOL> if last_word ( text_before_cursor , <EOL> '<STR_LIT>' ) . startswith ( '<STR_LIT:(>' ) : <EOL> return [ { '<STR_LIT:type>' : '<STR_LIT>' } ] <EOL> elif p . token_first ( ) . value . lower ( ) == '<STR_LIT>' : <EOL> return [ { '<STR_LIT:type>' : '<STR_LIT>' } ] <EOL> return [ { '<STR_LIT:type>' : '<STR_LIT>' , '<STR_LIT>' : extract_tables ( full_text ) } ] <EOL> elif token_v in ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ) : <EOL> return [ { '<STR_LIT:type>' : '<STR_LIT>' , '<STR_LIT>' : extract_tables ( full_text ) } ] <EOL> elif token_v in ( '<STR_LIT>' ) : <EOL> return [ { '<STR_LIT:type>' : '<STR_LIT>' } ] <EOL> elif token_v in ( '<STR_LIT:to>' , ) : <EOL> p = sqlparse . parse ( text_before_cursor ) [ <NUM_LIT:0> ] <EOL> if p . token_first ( ) . value . lower ( ) == '<STR_LIT>' : <EOL> return [ { '<STR_LIT:type>' : '<STR_LIT>' } ] <EOL> else : <EOL> return [ { '<STR_LIT:type>' : '<STR_LIT:user>' } ] <EOL> elif token_v in ( '<STR_LIT:user>' , '<STR_LIT>' ) : <EOL> return [ { '<STR_LIT:type>' : '<STR_LIT:user>' } ] <EOL> elif token_v in ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ) : <EOL> parent = ( identifier and identifier . get_parent_name ( ) ) or [ ] <EOL> if parent : <EOL> tables = extract_tables ( full_text ) <EOL> tables = [ t for t in tables if identifies ( parent , * t ) ] <EOL> return [ { '<STR_LIT:type>' : '<STR_LIT>' , '<STR_LIT>' : tables } , <EOL> { '<STR_LIT:type>' : '<STR_LIT>' , '<STR_LIT>' : parent } , <EOL> { '<STR_LIT:type>' : '<STR_LIT>' , '<STR_LIT>' : parent } , <EOL> { '<STR_LIT:type>' : '<STR_LIT>' , '<STR_LIT>' : parent } ] <EOL> else : <EOL> return [ { '<STR_LIT:type>' : '<STR_LIT>' , '<STR_LIT>' : extract_tables ( full_text ) } , <EOL> { '<STR_LIT:type>' : '<STR_LIT>' , '<STR_LIT>' : [ ] } , <EOL> { '<STR_LIT:type>' : '<STR_LIT>' } ] <EOL> elif ( token_v . endswith ( '<STR_LIT>' ) and token . is_keyword ) or ( token_v in <EOL> ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT>' ) ) : <EOL> schema = ( identifier and identifier . get_parent_name ( ) ) or [ ] <EOL> suggest = [ { '<STR_LIT:type>' : '<STR_LIT>' , '<STR_LIT>' : schema } ] <EOL> if not schema : <EOL> suggest . insert ( <NUM_LIT:0> , { '<STR_LIT:type>' : '<STR_LIT>' } ) <EOL> if token_v != '<STR_LIT>' : <EOL> suggest . append ( { '<STR_LIT:type>' : '<STR_LIT>' , '<STR_LIT>' : schema } ) <EOL> return suggest <EOL> elif token_v in ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ) : <EOL> rel_type = token_v <EOL> schema = ( identifier and identifier . get_parent_name ( ) ) or [ ] <EOL> if schema : <EOL> return [ { '<STR_LIT:type>' : rel_type , '<STR_LIT>' : schema } ] <EOL> else : <EOL> return [ { '<STR_LIT:type>' : '<STR_LIT>' } , { '<STR_LIT:type>' : rel_type , '<STR_LIT>' : [ ] } ] <EOL> elif token_v == '<STR_LIT>' : <EOL> tables = extract_tables ( full_text ) <EOL> parent = ( identifier and identifier . get_parent_name ( ) ) or [ ] <EOL> if parent : <EOL> tables = [ t for t in tables if identifies ( parent , * t ) ] <EOL> return [ { '<STR_LIT:type>' : '<STR_LIT>' , '<STR_LIT>' : tables } , <EOL> { '<STR_LIT:type>' : '<STR_LIT>' , '<STR_LIT>' : parent } , <EOL> { '<STR_LIT:type>' : '<STR_LIT>' , '<STR_LIT>' : parent } , <EOL> { '<STR_LIT:type>' : '<STR_LIT>' , '<STR_LIT>' : parent } ] <EOL> else : <EOL> aliases = [ t [ <NUM_LIT:2> ] or t [ <NUM_LIT:1> ] for t in tables ] <EOL> suggest = [ { '<STR_LIT:type>' : '<STR_LIT>' , '<STR_LIT>' : aliases } ] <EOL> if not aliases : <EOL> suggest . append ( { '<STR_LIT:type>' : '<STR_LIT>' , '<STR_LIT>' : parent } ) <EOL> return suggest <EOL> elif token_v in ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ) : <EOL> return [ { '<STR_LIT:type>' : '<STR_LIT>' } ] <EOL> elif token_v == '<STR_LIT>' : <EOL> return [ { '<STR_LIT:type>' : '<STR_LIT>' } ] <EOL> elif token_v . endswith ( '<STR_LIT:U+002C>' ) or token_v in [ '<STR_LIT:=>' , '<STR_LIT>' , '<STR_LIT>' ] : <EOL> prev_keyword , text_before_cursor = find_prev_keyword ( text_before_cursor ) <EOL> if prev_keyword : <EOL> return suggest_based_on_last_token ( <EOL> prev_keyword , text_before_cursor , full_text , identifier ) <EOL> else : <EOL> return [ ] <EOL> else : <EOL> return [ { '<STR_LIT:type>' : '<STR_LIT>' } ] <EOL> def identifies ( id , schema , table , alias ) : <EOL> return id == alias or id == table or ( <EOL> schema and ( id == schema + '<STR_LIT:.>' + table ) ) </s>
