Merge branch 'type-hint' of git://github.com/letientai299/vim-pydocstring

Author: heavenshell

autoload/pydocstring.vim

@@ -24,10 +24,10 @@ if !exists('g:pydocstring_ignore_args_pattern')
 endif
 
 let s:regexs = {
-\ 'def': '^def\s\|^\s*def\s',
-\ 'class': '^class\s\|^\s*class\s',
-\ 'async': '^async\s*def\s\|^\s*async\sdef\s'
-\ }
+      \ 'def': '^def\s\|^\s*def\s',
+      \ 'class': '^class\s\|^\s*class\s',
+      \ 'async': '^async\s*def\s\|^\s*async\sdef\s'
+      \ }
 
 function! s:readtmpl(type)
   let tmpldir = g:pydocstring_templates_dir
@@ -44,33 +44,54 @@ function! s:readtmpl(type)
   return tmpl
 endfunction
 
+function! s:parseClass(line)
+  " For class definition, we just simply need to extract the class name.  We can
+  " do that by just delete every white spaces and the whole parenthesics if
+  " existed.
+  let header = substitute(a:line, '\s\|(.*\|:', '', 'g')
+  let parse = {'type': 'class', 'header': header, 'args': '', 'returnType': ''}
+  return parse
+endfunction
+
+function! s:parseFunc(type, line)
+  let header = substitute(a:line, '\s\|(.*\|:', '', 'g')
+
+  let argsStr = substitute(a:line, '\s\|.*(\|).*', '', 'g')
+  let args = split(argsStr, ',')
+
+  let arrowIndex = match(a:line, "->")
+  if arrowIndex != -1
+    let substring = strpart(a:line, arrowIndex+2)
+    let returnType = substitute(substring, '\W*', '', 'g')
+  else
+    let returnType = ''
+  endif
+
+  let parse = {'type': a:type, 'header': header, 'args': args, 'returnType': returnType}
+  return parse
+endfunction
+
 function! s:parse(line)
   let str = substitute(a:line, '\\', '', 'g')
+  let str = substitute(a:line, '#.*$', '', 'g')
   let type = ''
+
+  if str =~ s:regexs['class']
+    let str = substitute(str, s:regexs['class'], '', '')
+    return s:parseClass(str)
+  endif
+
   if str =~ s:regexs['def']
     let str = substitute(str, s:regexs['def'], '', '')
     let type = 'def'
   elseif str =~ s:regexs['async']
     let str = substitute(str, s:regexs['async'], '', '')
     let type = 'def'
-  elseif str =~ s:regexs['class']
-    let str = substitute(str, s:regexs['class'], '', '')
-    let type = 'class'
   else
     return 0
   endif
-  let str = substitute(str, '\s\|):\|)\s:', '', 'g')
 
-  let strs = split(str, '(')
-  let header = strs[0]
-  let args = []
-  if len(strs) > 1
-    let args = split(strs[1], ',')
-  end
-
-  let parse = {'type': type, 'header': header, 'args': args}
-
-  return parse
+  return s:parseFunc(type, str)
 endfunction
 
 " Vim Script does not support lambda function...
@@ -81,74 +102,101 @@ function! s:readoneline(indent, prefix)
   return tmpl
 endfunction
 
+" Check if we should show args in the docstring. We won't do that in  case:
+" - There's no args.
+" - There's only one arg that match with g:pydocstring_ignore_args_pattern
+function! s:shouldIncludeArgs(args)
+  if len(a:args) == 0
+    return 0
+  endif
+
+  if len(a:args) == 1 && a:args[0] =~ g:pydocstring_ignore_args_pattern
+    return 0
+  endif
+
+  return 1
+endfunction
+
+" Check if we should use one line docstring.
+" There's several cases:
+" - Type is `class`
+" - No return type and no args.
+" - No return type and the only one args is `self` or `cls` (defined by
+"   g:pydocstring_ignore_args_pattern
+"
+" Return 1 for True, and 0 for False
+function! s:shouldUseOneLineDocString(type, args, returnType)
+  if a:type != 'def'
+    return 1
+  endif
+
+  if a:returnType != ''
+    return 0
+  endif
+
+  return !s:shouldIncludeArgs(a:args)
+endfunction
+
 function! s:builddocstring(strs, indent, nested_indent)
   let type  = a:strs['type']
   let prefix = a:strs['header']
   let args = a:strs['args']
+  let returnType = a:strs['returnType']
+
+  if s:shouldUseOneLineDocString(type, args, returnType)
+    return s:readoneline(a:indent, prefix)
+  endif
+
   let tmpl = ''
-  if len(args) > 0 && type == 'def'
-    let docstrings = []
-    let lines = s:readtmpl('multi')
-    for line in lines
-      if line =~ '{{_header_}}'
-        let header = substitute(line, '{{_header_}}', prefix, '')
-        call add(docstrings, a:indent . header)
-      elseif line =~ '{{_arg_}}'
-        if len(args) == 0
-          let tmpl = s:readoneline(a:indent, prefix)
-          return tmpl
-        endif
-
-        if args[0] =~ g:pydocstring_ignore_args_pattern && len(args) == 1
-          let tmpl = s:readoneline(a:indent, prefix)
-          return tmpl
-        endif
-
-        let arglist = []
+  let docstrings = []
+  let lines = s:readtmpl('multi')
+  for line in lines
+    if line =~ '{{_header_}}'
+      let header = substitute(line, '{{_header_}}', prefix, '')
+      call add(docstrings, a:indent . header)
+      call add(docstrings, '' )
+    elseif line =~ '{{_args_}}'
+      if len(args) != 0
         for arg in args
           let arg = substitute(arg, '=.*$', '', '')
           if arg =~ g:pydocstring_ignore_args_pattern
             continue
           endif
-          let arg = substitute(line, '{{_arg_}}', arg, 'g')
-          let arg = substitute(arg, '{{_lf_}}', "\n", '')
-          let arg = substitute(arg, '{{_indent_}}', a:indent, 'g')
-          let arg = substitute(arg, '{{_nested_indent_}}', a:nested_indent, 'g')
-          call add(docstrings, a:indent . arg)
-        endfor
-      elseif line =~ '{{_indent_}}'
-        let arg = substitute(line, '{{_indent_}}', a:indent, 'g')
-        call add(docstrings, arg)
-      elseif line =~ '{{_args_}}'
-        if len(args) == 0
-          let tmpl = s:readoneline(a:indent, prefix)
-          return tmpl
-        endif
-
-        if args[0] =~ g:pydocstring_ignore_args_pattern && len(args) == 1
-          let tmpl = s:readoneline(a:indent, prefix)
-          return tmpl
-        endif
-
-        let arglist = []
-        for arg in args
-          let arg = substitute(arg, '=.*$', '', '')
-          if arg =~ g:pydocstring_ignore_args_pattern
-            continue
+          let template = line
+
+          if match(arg, ':') != -1
+            let argTemplate = s:readtmpl('arg')
+            let argTemplate = join(s:readtmpl('arg'), '')
+            let argParts = split(arg, ':')
+            let argTemplate = substitute(argTemplate, '{{_name_}}', argParts[0], '')
+            let arg = substitute(argTemplate, '{{_type_}}', argParts[1], '')
           endif
-          let arg = substitute(line, '{{_args_}}', arg, '')
-          call add(docstrings, a:indent . arg)
+
+          let template = substitute(template, '{{_args_}}', arg, 'g')
+          let template = substitute(template, '{{_lf_}}', '\n', '')
+          let template = substitute(template, '{{_indent_}}', a:indent, 'g')
+          let template = substitute(template, '{{_nested_indent_}}', a:nested_indent, 'g')
+          call add(docstrings, a:indent . template)
         endfor
-      elseif line == '"""'
-        call add(docstrings, a:indent . line)
+        call add(docstrings, '' )
+      endif
+    elseif line =~ '{{_indent_}}'
+      let arg = substitute(line, '{{_indent_}}', a:indent, 'g')
+      call add(docstrings, arg)
+    elseif line =~ '{{_returnType_}}'
+      if strlen(returnType) != 0
+        let rt = substitute(line, '{{_returnType_}}', returnType, '')
+        call add(docstrings, a:indent . rt)
       else
-        call add(docstrings, line)
+        call remove(docstrings, -1)
       endif
-    endfor
-    let tmpl = substitute(join(docstrings, "\n"), "\n$", '', '')
-  else
-    let tmpl = s:readoneline(a:indent, prefix)
-  endif
+    elseif line == '"""'
+      call add(docstrings, a:indent . line)
+    else
+      call add(docstrings, line)
+    endif
+  endfor
+  let tmpl = substitute(join(docstrings, "\n"), "\n$", '', '')
 
   return tmpl
 endfunction
@@ -159,7 +207,7 @@ function! pydocstring#insert()
   let indent = matchstr(line, '^\(\s*\)')
 
   let startpos = line('.')
-  let insertpos = search('\:\+$')
+  let insertpos = search('\:\(\s*#.*\)*$')
   let lines = join(getline(startpos, insertpos))
 
   let docstring = s:parse(lines)

doc/pydocstring.txt

@@ -124,23 +124,33 @@ TEMPLATE					*pydocstring-vim-template*
 
 If you don't like default docstring, You can modify docstring template.
 >
-  let g:template_vim_template_dir = '/path/to/your/template/directory'
+  let g:pydocstring_templates_dir = '/path/to/your/template/directory'
 <
 - comment.txt
   Comment.
 - oneline.txt
   One-line docstring.
 - multi.txt
   Multi-line docstring.
+- arg.txt
+  Template for each argument.
 
 Template variables.
   |{{_header_}}|	Class, function or method name as a value.
-  |{{_args_}}|	Assign function or method's arguments.
-  |{{_arg_}}|	Assign function or method's argument.
-  |{{_ln_}}|	Assign line break.
+  |{{_args_}}|		Assign function or method's arguments.
+  |{{_lf_}}|		Assign line break.
   |{{_indent_}}|	Assign indent.
   |{{_nested_indent_}}|	Assign indent.
-
+  |{{_returnType_}}|	Function or method return type.
+  |{{_name_}}|		Argument name, should only be used in `arg.txt`
+  |{{_type_}}|		Argument type, should only be used in `arg.txt`
+
+There's some rules for parsing template:
+- `{{_header_}`, `{{_args}}` and `{{_returnType_}}` should be on their own
+  line.
+- Pydocstring will insert a emmpty line between each parts, so you don't
+  need to include empty line in your template. Please see the default template
+  directory to have an example.
 ==============================================================================
 VARIABLES					*pydocstring-vim-variables*
 

template/pydocstring/arg.txt

@@ -0,0 +1 @@
+{{_name_}} ({{_type_}})

template/pydocstring/multi.txt

@@ -1,4 +1,4 @@
 """{{_header_}}
-
 :param {{_args_}}:
+return ({{_returnType_}}):
 """

test/basic.vader

@@ -45,6 +45,23 @@ Expect python:
         """
         pass
 
+Given python (def foo with 2 params and a comment block):
+    def foo(arg1, arg2): # Something fun (just, for, test)
+        pass
+
+Execute:
+    Pydocstring
+
+Expect python:
+    def foo(arg1, arg2): # Something fun (just, for, test)
+        """foo
+
+        :param arg1:
+        :param arg2:
+        """
+        pass
+
+
 Given python (def foo with variadic params):
     def foo(*arg):
         pass
@@ -64,6 +81,31 @@ Expect python:
 #-------------------------------------------------------------------------------
 
 
+Given python (class Foo no extends):
+    class Foo:
+        pass
+
+Execute:
+    Pydocstring
+
+Expect python:
+    class Foo:
+        """Foo"""
+        pass
+
+Given python (class Foo extends multiple classes):
+    class Foo(ClassA, ClassB, ClassC):
+        pass
+
+Execute:
+    Pydocstring
+
+Expect python:
+    class Foo(ClassA, ClassB, ClassC):
+        """Foo"""
+        pass
+
+
 Given python (class Foo):
     class Foo(object):
         def foo(self):

test/custom-template.vader

@@ -29,7 +29,7 @@ Expect python:
         """foo
 
         arg1:
-            arg1 is
+            the ...
         """
         pass
 
@@ -46,9 +46,9 @@ Expect python:
         """foo
 
         arg1:
-            arg1 is
+            the ...
         arg2:
-            arg2 is
+            the ...
         """
         pass
 

test/nodemon.json

@@ -6,7 +6,7 @@
     "vader.vim"
   ],
   "verbose": true,
-  "ext": "vader vim",
+  "ext": "txt vader vim",
   "execMap": {
     "vader": "./run-single-test-file.sh"
   }

test/test-template/arg.txt

@@ -0,0 +1 @@
+{{_type_}} {{_name_}}

test/test-template/multi.txt

@@ -1,4 +1,4 @@
 """{{_header_}}
-
-{{_arg_}}:{{_lf_}}{{_indent_}}    {{_arg_}} is
+{{_args_}}:{{_lf_}}{{_indent_}}    the ...
+return ({{_returnType_}}):
 """