Support type hint for arguments

Author: letientai299

autoload/pydocstring.vim

@@ -163,6 +163,15 @@ function! s:builddocstring(strs, indent, nested_indent)
             continue
           endif
           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 template = substitute(template, '{{_args_}}', arg, 'g')
           let template = substitute(template, '{{_lf_}}', '\n', '')
           let template = substitute(template, '{{_indent_}}', a:indent, 'g')

doc/pydocstring.txt

@@ -132,6 +132,8 @@ If you don't like default docstring, You can modify docstring template.
   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.
@@ -140,6 +142,8 @@ Template variables.
   |{{_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

template/pydocstring/arg.txt

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

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/test-template/arg.txt

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

test/test-template/multi.txt

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

test/type-hint-custom-template.vader

@@ -0,0 +1,27 @@
+# vim:set et sw=4 ts=4 tw=79:
+Execute (Setup template dir):
+    Save g:pydocstring_templates_dir
+    let g:pydocstring_templates_dir = './test-template/'
+
+Given python (def foo 1 arg with type, 1 none, and return int):
+    def foo(n: int, arg) -> int:
+        pass
+
+Execute:
+    Pydocstring
+
+Expect python:
+    def foo(n: int, arg) -> int:
+        """foo
+
+        int n:
+            the ...
+        arg:
+            the ...
+
+        return (int):
+        """
+        pass
+
+Execute (Clear pydocstring_templates_dir):
+    Restore

test/type-hint.vader

@@ -3,6 +3,7 @@
 Given python (def foo no arg return str):
     def foo() -> str:
         pass
+
 Execute:
     Pydocstring
 
@@ -46,3 +47,58 @@ Expect python:
         return (str):
         """
         pass
+
+Given python (def foo 1 arg with type and return int):
+    def foo(n: int) -> int:
+        pass
+
+Execute:
+    Pydocstring
+
+Expect python:
+    def foo(n: int) -> int:
+        """foo
+
+        :param n (int):
+
+        return (int):
+        """
+        pass
+
+Given python (def foo 2 args with type and return int):
+    def foo(n: int, m: str) -> int:
+        pass
+
+Execute:
+    Pydocstring
+
+Expect python:
+    def foo(n: int, m: str) -> int:
+        """foo
+
+        :param n (int):
+        :param m (str):
+
+        return (int):
+        """
+        pass
+
+
+Given python (def foo 1 arg with type, 1 none, and return int):
+    def foo(n: int, arg) -> int:
+        pass
+
+Execute:
+    Pydocstring
+
+Expect python:
+    def foo(n: int, arg) -> int:
+        """foo
+
+        :param n (int):
+        :param arg:
+
+        return (int):
+        """
+        pass
+