1"pythoncomplete.vim - Omni Completion for python
2" Maintainer: Aaron Griffin <aaronmgriffin@gmail.com>
3" Version: 0.9
4" Last Updated: 18 Jun 2009
5"
6" Changes
7" TODO:
8" 'info' item output can use some formatting work
9" Add an "unsafe eval" mode, to allow for return type evaluation
10" Complete basic syntax along with import statements
11"   i.e. "import url<c-x,c-o>"
12" Continue parsing on invalid line??
13"
14" v 0.9
15"   * Fixed docstring parsing for classes and functions
16"   * Fixed parsing of *args and **kwargs type arguments
17"   * Better function param parsing to handle things like tuples and
18"     lambda defaults args
19"
20" v 0.8
21"   * Fixed an issue where the FIRST assignment was always used instead of
22"   using a subsequent assignment for a variable
23"   * Fixed a scoping issue when working inside a parameterless function
24"
25"
26" v 0.7
27"   * Fixed function list sorting (_ and __ at the bottom)
28"   * Removed newline removal from docs.  It appears vim handles these better in
29"   recent patches
30"
31" v 0.6:
32"   * Fixed argument completion
33"   * Removed the 'kind' completions, as they are better indicated
34"   with real syntax
35"   * Added tuple assignment parsing (whoops, that was forgotten)
36"   * Fixed import handling when flattening scope
37"
38" v 0.5:
39" Yeah, I skipped a version number - 0.4 was never public.
40"  It was a bugfix version on top of 0.3.  This is a complete
41"  rewrite.
42"
43
44if !has('python')
45    echo "Error: Required vim compiled with +python"
46    finish
47endif
48
49function! pythoncomplete#Complete(findstart, base)
50    "findstart = 1 when we need to get the text length
51    if a:findstart == 1
52        let line = getline('.')
53        let idx = col('.')
54        while idx > 0
55            let idx -= 1
56            let c = line[idx]
57            if c =~ '\w'
58                continue
59            elseif ! c =~ '\.'
60                let idx = -1
61                break
62            else
63                break
64            endif
65        endwhile
66
67        return idx
68    "findstart = 0 when we need to return the list of completions
69    else
70        "vim no longer moves the cursor upon completion... fix that
71        let line = getline('.')
72        let idx = col('.')
73        let cword = ''
74        while idx > 0
75            let idx -= 1
76            let c = line[idx]
77            if c =~ '\w' || c =~ '\.'
78                let cword = c . cword
79                continue
80            elseif strlen(cword) > 0 || idx == 0
81                break
82            endif
83        endwhile
84        execute "python vimcomplete('" . cword . "', '" . a:base . "')"
85        return g:pythoncomplete_completions
86    endif
87endfunction
88
89function! s:DefPython()
90python << PYTHONEOF
91import sys, tokenize, cStringIO, types
92from token import NAME, DEDENT, NEWLINE, STRING
93
94debugstmts=[]
95def dbg(s): debugstmts.append(s)
96def showdbg():
97    for d in debugstmts: print "DBG: %s " % d
98
99def vimcomplete(context,match):
100    global debugstmts
101    debugstmts = []
102    try:
103        import vim
104        def complsort(x,y):
105            try:
106                xa = x['abbr']
107                ya = y['abbr']
108                if xa[0] == '_':
109                    if xa[1] == '_' and ya[0:2] == '__':
110                        return xa > ya
111                    elif ya[0:2] == '__':
112                        return -1
113                    elif y[0] == '_':
114                        return xa > ya
115                    else:
116                        return 1
117                elif ya[0] == '_':
118                    return -1
119                else:
120                   return xa > ya
121            except:
122                return 0
123        cmpl = Completer()
124        cmpl.evalsource('\n'.join(vim.current.buffer),vim.eval("line('.')"))
125        all = cmpl.get_completions(context,match)
126        all.sort(complsort)
127        dictstr = '['
128        # have to do this for double quoting
129        for cmpl in all:
130            dictstr += '{'
131            for x in cmpl: dictstr += '"%s":"%s",' % (x,cmpl[x])
132            dictstr += '"icase":0},'
133        if dictstr[-1] == ',': dictstr = dictstr[:-1]
134        dictstr += ']'
135        #dbg("dict: %s" % dictstr)
136        vim.command("silent let g:pythoncomplete_completions = %s" % dictstr)
137        #dbg("Completion dict:\n%s" % all)
138    except vim.error:
139        dbg("VIM Error: %s" % vim.error)
140
141class Completer(object):
142    def __init__(self):
143       self.compldict = {}
144       self.parser = PyParser()
145
146    def evalsource(self,text,line=0):
147        sc = self.parser.parse(text,line)
148        src = sc.get_code()
149        dbg("source: %s" % src)
150        try: exec(src) in self.compldict
151        except: dbg("parser: %s, %s" % (sys.exc_info()[0],sys.exc_info()[1]))
152        for l in sc.locals:
153            try: exec(l) in self.compldict
154            except: dbg("locals: %s, %s [%s]" % (sys.exc_info()[0],sys.exc_info()[1],l))
155
156    def _cleanstr(self,doc):
157        return doc.replace('"',' ').replace("'",' ')
158
159    def get_arguments(self,func_obj):
160        def _ctor(obj):
161            try: return class_ob.__init__.im_func
162            except AttributeError:
163                for base in class_ob.__bases__:
164                    rc = _find_constructor(base)
165                    if rc is not None: return rc
166            return None
167
168        arg_offset = 1
169        if type(func_obj) == types.ClassType: func_obj = _ctor(func_obj)
170        elif type(func_obj) == types.MethodType: func_obj = func_obj.im_func
171        else: arg_offset = 0
172        
173        arg_text=''
174        if type(func_obj) in [types.FunctionType, types.LambdaType]:
175            try:
176                cd = func_obj.func_code
177                real_args = cd.co_varnames[arg_offset:cd.co_argcount]
178                defaults = func_obj.func_defaults or ''
179                defaults = map(lambda name: "=%s" % name, defaults)
180                defaults = [""] * (len(real_args)-len(defaults)) + defaults
181                items = map(lambda a,d: a+d, real_args, defaults)
182                if func_obj.func_code.co_flags & 0x4:
183                    items.append("...")
184                if func_obj.func_code.co_flags & 0x8:
185                    items.append("***")
186                arg_text = (','.join(items)) + ')'
187
188            except:
189                dbg("arg completion: %s: %s" % (sys.exc_info()[0],sys.exc_info()[1]))
190                pass
191        if len(arg_text) == 0:
192            # The doc string sometimes contains the function signature
193            #  this works for alot of C modules that are part of the
194            #  standard library
195            doc = func_obj.__doc__
196            if doc:
197                doc = doc.lstrip()
198                pos = doc.find('\n')
199                if pos > 0:
200                    sigline = doc[:pos]
201                    lidx = sigline.find('(')
202                    ridx = sigline.find(')')
203                    if lidx > 0 and ridx > 0:
204                        arg_text = sigline[lidx+1:ridx] + ')'
205        if len(arg_text) == 0: arg_text = ')'
206        return arg_text
207
208    def get_completions(self,context,match):
209        dbg("get_completions('%s','%s')" % (context,match))
210        stmt = ''
211        if context: stmt += str(context)
212        if match: stmt += str(match)
213        try:
214            result = None
215            all = {}
216            ridx = stmt.rfind('.')
217            if len(stmt) > 0 and stmt[-1] == '(':
218                result = eval(_sanitize(stmt[:-1]), self.compldict)
219                doc = result.__doc__
220                if doc is None: doc = ''
221                args = self.get_arguments(result)
222                return [{'word':self._cleanstr(args),'info':self._cleanstr(doc)}]
223            elif ridx == -1:
224                match = stmt
225                all = self.compldict
226            else:
227                match = stmt[ridx+1:]
228                stmt = _sanitize(stmt[:ridx])
229                result = eval(stmt, self.compldict)
230                all = dir(result)
231
232            dbg("completing: stmt:%s" % stmt)
233            completions = []
234
235            try: maindoc = result.__doc__
236            except: maindoc = ' '
237            if maindoc is None: maindoc = ' '
238            for m in all:
239                if m == "_PyCmplNoType": continue #this is internal
240                try:
241                    dbg('possible completion: %s' % m)
242                    if m.find(match) == 0:
243                        if result is None: inst = all[m]
244                        else: inst = getattr(result,m)
245                        try: doc = inst.__doc__
246                        except: doc = maindoc
247                        typestr = str(inst)
248                        if doc is None or doc == '': doc = maindoc
249
250                        wrd = m[len(match):]
251                        c = {'word':wrd, 'abbr':m,  'info':self._cleanstr(doc)}
252                        if "function" in typestr:
253                            c['word'] += '('
254                            c['abbr'] += '(' + self._cleanstr(self.get_arguments(inst))
255                        elif "method" in typestr:
256                            c['word'] += '('
257                            c['abbr'] += '(' + self._cleanstr(self.get_arguments(inst))
258                        elif "module" in typestr:
259                            c['word'] += '.'
260                        elif "class" in typestr:
261                            c['word'] += '('
262                            c['abbr'] += '('
263                        completions.append(c)
264                except:
265                    i = sys.exc_info()
266                    dbg("inner completion: %s,%s [stmt='%s']" % (i[0],i[1],stmt))
267            return completions
268        except:
269            i = sys.exc_info()
270            dbg("completion: %s,%s [stmt='%s']" % (i[0],i[1],stmt))
271            return []
272
273class Scope(object):
274    def __init__(self,name,indent,docstr=''):
275        self.subscopes = []
276        self.docstr = docstr
277        self.locals = []
278        self.parent = None
279        self.name = name
280        self.indent = indent
281
282    def add(self,sub):
283        #print 'push scope: [%s@%s]' % (sub.name,sub.indent)
284        sub.parent = self
285        self.subscopes.append(sub)
286        return sub
287
288    def doc(self,str):
289        """ Clean up a docstring """
290        d = str.replace('\n',' ')
291        d = d.replace('\t',' ')
292        while d.find('  ') > -1: d = d.replace('  ',' ')
293        while d[0] in '"\'\t ': d = d[1:]
294        while d[-1] in '"\'\t ': d = d[:-1]
295        dbg("Scope(%s)::docstr = %s" % (self,d))
296        self.docstr = d
297
298    def local(self,loc):
299        self._checkexisting(loc)
300        self.locals.append(loc)
301
302    def copy_decl(self,indent=0):
303        """ Copy a scope's declaration only, at the specified indent level - not local variables """
304        return Scope(self.name,indent,self.docstr)
305
306    def _checkexisting(self,test):
307        "Convienance function... keep out duplicates"
308        if test.find('=') > -1:
309            var = test.split('=')[0].strip()
310            for l in self.locals:
311                if l.find('=') > -1 and var == l.split('=')[0].strip():
312                    self.locals.remove(l)
313
314    def get_code(self):
315        str = ""
316        if len(self.docstr) > 0: str += '"""'+self.docstr+'"""\n'
317        for l in self.locals:
318            if l.startswith('import'): str += l+'\n'
319        str += 'class _PyCmplNoType:\n    def __getattr__(self,name):\n        return None\n'
320        for sub in self.subscopes:
321            str += sub.get_code()
322        for l in self.locals:
323            if not l.startswith('import'): str += l+'\n'
324
325        return str
326
327    def pop(self,indent):
328        #print 'pop scope: [%s] to [%s]' % (self.indent,indent)
329        outer = self
330        while outer.parent != None and outer.indent >= indent:
331            outer = outer.parent
332        return outer
333
334    def currentindent(self):
335        #print 'parse current indent: %s' % self.indent
336        return '    '*self.indent
337
338    def childindent(self):
339        #print 'parse child indent: [%s]' % (self.indent+1)
340        return '    '*(self.indent+1)
341
342class Class(Scope):
343    def __init__(self, name, supers, indent, docstr=''):
344        Scope.__init__(self,name,indent, docstr)
345        self.supers = supers
346    def copy_decl(self,indent=0):
347        c = Class(self.name,self.supers,indent, self.docstr)
348        for s in self.subscopes:
349            c.add(s.copy_decl(indent+1))
350        return c
351    def get_code(self):
352        str = '%sclass %s' % (self.currentindent(),self.name)
353        if len(self.supers) > 0: str += '(%s)' % ','.join(self.supers)
354        str += ':\n'
355        if len(self.docstr) > 0: str += self.childindent()+'"""'+self.docstr+'"""\n'
356        if len(self.subscopes) > 0:
357            for s in self.subscopes: str += s.get_code()
358        else:
359            str += '%spass\n' % self.childindent()
360        return str
361
362
363class Function(Scope):
364    def __init__(self, name, params, indent, docstr=''):
365        Scope.__init__(self,name,indent, docstr)
366        self.params = params
367    def copy_decl(self,indent=0):
368        return Function(self.name,self.params,indent, self.docstr)
369    def get_code(self):
370        str = "%sdef %s(%s):\n" % \
371            (self.currentindent(),self.name,','.join(self.params))
372        if len(self.docstr) > 0: str += self.childindent()+'"""'+self.docstr+'"""\n'
373        str += "%spass\n" % self.childindent()
374        return str
375
376class PyParser:
377    def __init__(self):
378        self.top = Scope('global',0)
379        self.scope = self.top
380
381    def _parsedotname(self,pre=None):
382        #returns (dottedname, nexttoken)
383        name = []
384        if pre is None:
385            tokentype, token, indent = self.next()
386            if tokentype != NAME and token != '*':
387                return ('', token)
388        else: token = pre
389        name.append(token)
390        while True:
391            tokentype, token, indent = self.next()
392            if token != '.': break
393            tokentype, token, indent = self.next()
394            if tokentype != NAME: break
395            name.append(token)
396        return (".".join(name), token)
397
398    def _parseimportlist(self):
399        imports = []
400        while True:
401            name, token = self._parsedotname()
402            if not name: break
403            name2 = ''
404            if token == 'as': name2, token = self._parsedotname()
405            imports.append((name, name2))
406            while token != "," and "\n" not in token:
407                tokentype, token, indent = self.next()
408            if token != ",": break
409        return imports
410
411    def _parenparse(self):
412        name = ''
413        names = []
414        level = 1
415        while True:
416            tokentype, token, indent = self.next()
417            if token in (')', ',') and level == 1:
418                if '=' not in name: name = name.replace(' ', '')
419                names.append(name.strip())
420                name = ''
421            if token == '(':
422                level += 1
423                name += "("
424            elif token == ')':
425                level -= 1
426                if level == 0: break
427                else: name += ")"
428            elif token == ',' and level == 1:
429                pass
430            else:
431                name += "%s " % str(token)
432        return names
433
434    def _parsefunction(self,indent):
435        self.scope=self.scope.pop(indent)
436        tokentype, fname, ind = self.next()
437        if tokentype != NAME: return None
438
439        tokentype, open, ind = self.next()
440        if open != '(': return None
441        params=self._parenparse()
442
443        tokentype, colon, ind = self.next()
444        if colon != ':': return None
445
446        return Function(fname,params,indent)
447
448    def _parseclass(self,indent):
449        self.scope=self.scope.pop(indent)
450        tokentype, cname, ind = self.next()
451        if tokentype != NAME: return None
452
453        super = []
454        tokentype, next, ind = self.next()
455        if next == '(':
456            super=self._parenparse()
457        elif next != ':': return None
458
459        return Class(cname,super,indent)
460
461    def _parseassignment(self):
462        assign=''
463        tokentype, token, indent = self.next()
464        if tokentype == tokenize.STRING or token == 'str':  
465            return '""'
466        elif token == '(' or token == 'tuple':
467            return '()'
468        elif token == '[' or token == 'list':
469            return '[]'
470        elif token == '{' or token == 'dict':
471            return '{}'
472        elif tokentype == tokenize.NUMBER:
473            return '0'
474        elif token == 'open' or token == 'file':
475            return 'file'
476        elif token == 'None':
477            return '_PyCmplNoType()'
478        elif token == 'type':
479            return 'type(_PyCmplNoType)' #only for method resolution
480        else:
481            assign += token
482            level = 0
483            while True:
484                tokentype, token, indent = self.next()
485                if token in ('(','{','['):
486                    level += 1
487                elif token in (']','}',')'):
488                    level -= 1
489                    if level == 0: break
490                elif level == 0:
491                    if token in (';','\n'): break
492                    assign += token
493        return "%s" % assign
494
495    def next(self):
496        type, token, (lineno, indent), end, self.parserline = self.gen.next()
497        if lineno == self.curline:
498            #print 'line found [%s] scope=%s' % (line.replace('\n',''),self.scope.name)
499            self.currentscope = self.scope
500        return (type, token, indent)
501
502    def _adjustvisibility(self):
503        newscope = Scope('result',0)
504        scp = self.currentscope
505        while scp != None:
506            if type(scp) == Function:
507                slice = 0
508                #Handle 'self' params
509                if scp.parent != None and type(scp.parent) == Class:
510                    slice = 1
511                    newscope.local('%s = %s' % (scp.params[0],scp.parent.name))
512                for p in scp.params[slice:]:
513                    i = p.find('=')
514                    if len(p) == 0: continue
515                    pvar = ''
516                    ptype = ''
517                    if i == -1:
518                        pvar = p
519                        ptype = '_PyCmplNoType()'
520                    else:
521                        pvar = p[:i]
522                        ptype = _sanitize(p[i+1:])
523                    if pvar.startswith('**'):
524                        pvar = pvar[2:]
525                        ptype = '{}'
526                    elif pvar.startswith('*'):
527                        pvar = pvar[1:]
528                        ptype = '[]'
529
530                    newscope.local('%s = %s' % (pvar,ptype))
531
532            for s in scp.subscopes:
533                ns = s.copy_decl(0)
534                newscope.add(ns)
535            for l in scp.locals: newscope.local(l)
536            scp = scp.parent
537
538        self.currentscope = newscope
539        return self.currentscope
540
541    #p.parse(vim.current.buffer[:],vim.eval("line('.')"))
542    def parse(self,text,curline=0):
543        self.curline = int(curline)
544        buf = cStringIO.StringIO(''.join(text) + '\n')
545        self.gen = tokenize.generate_tokens(buf.readline)
546        self.currentscope = self.scope
547
548        try:
549            freshscope=True
550            while True:
551                tokentype, token, indent = self.next()
552                #dbg( 'main: token=[%s] indent=[%s]' % (token,indent))
553
554                if tokentype == DEDENT or token == "pass":
555                    self.scope = self.scope.pop(indent)
556                elif token == 'def':
557                    func = self._parsefunction(indent)
558                    if func is None:
559                        print "function: syntax error..."
560                        continue
561                    dbg("new scope: function")
562                    freshscope = True
563                    self.scope = self.scope.add(func)
564                elif token == 'class':
565                    cls = self._parseclass(indent)
566                    if cls is None:
567                        print "class: syntax error..."
568                        continue
569                    freshscope = True
570                    dbg("new scope: class")
571                    self.scope = self.scope.add(cls)
572                    
573                elif token == 'import':
574                    imports = self._parseimportlist()
575                    for mod, alias in imports:
576                        loc = "import %s" % mod
577                        if len(alias) > 0: loc += " as %s" % alias
578                        self.scope.local(loc)
579                    freshscope = False
580                elif token == 'from':
581                    mod, token = self._parsedotname()
582                    if not mod or token != "import":
583                        print "from: syntax error..."
584                        continue
585                    names = self._parseimportlist()
586                    for name, alias in names:
587                        loc = "from %s import %s" % (mod,name)
588                        if len(alias) > 0: loc += " as %s" % alias
589                        self.scope.local(loc)
590                    freshscope = False
591                elif tokentype == STRING:
592                    if freshscope: self.scope.doc(token)
593                elif tokentype == NAME:
594                    name,token = self._parsedotname(token) 
595                    if token == '=':
596                        stmt = self._parseassignment()
597                        dbg("parseassignment: %s = %s" % (name, stmt))
598                        if stmt != None:
599                            self.scope.local("%s = %s" % (name,stmt))
600                    freshscope = False
601        except StopIteration: #thrown on EOF
602            pass
603        except:
604            dbg("parse error: %s, %s @ %s" %
605                (sys.exc_info()[0], sys.exc_info()[1], self.parserline))
606        return self._adjustvisibility()
607
608def _sanitize(str):
609    val = ''
610    level = 0
611    for c in str:
612        if c in ('(','{','['):
613            level += 1
614        elif c in (']','}',')'):
615            level -= 1
616        elif level == 0:
617            val += c
618    return val
619
620sys.path.extend(['.','..'])
621PYTHONEOF
622endfunction
623
624call s:DefPython()
625" vim: set et ts=4:
626