Merge pull request #19 from pixelneo/hints

Support type hints

Author: pixelneo

README.md

@@ -1,10 +1,10 @@
 # vim-python-docstring
-This is a plugin to Vim for creating of docstrings. 
+This is a plugin to Vim and NeoVim for creating of docstrings. 
 
-Works also with neovim.
+**New**: Support for type hints and async functions.
 
 ## What it does
-Docstrings for methods will contain a **list of arguments**, **list of raised exceptions** and whether the method **yields** or **raises**.
+Docstrings for methods will contain a **list of parameters and their type hints**, **list of raised exceptions** and whether the method **yields** or **raises**.
 
 Class docstring will have a **list of atributes**. 
 
@@ -39,6 +39,7 @@ The plugin uses these commands:
 | Command       | Description |
 |---------------|-------------|
 | Docstring     | Create full docstring 
+| DocstringTypes| Just like `:Docstring` but includes type hints
 | DocstringLine | Create empty one-line docstring  
 
 ## Options:
@@ -64,7 +65,7 @@ Possible values = [`'google'`, `'numpy'`, `'rest'`, `'epytext'`]
 let g:python_style = 'google'
 ~~~
 
-## Work in progress
+## Development
 Pull requests are welcome as are feature request and issue reports.
 
 You can encounter some situations in which the plugin may not work as expected.

autoload/vimpythondocstring.vim

@@ -16,6 +16,10 @@ function! vimpythondocstring#Full()
     python3 pydocstring.Docstring().full_docstring()
 endfunction
 
+function! vimpythondocstring#FullTypes()
+    python3 pydocstring.Docstring().full_docstring(print_hints=True)
+endfunction
+
 function! vimpythondocstring#Oneline()
     python3 pydocstring.Docstring().oneline_docstring()
 endfunction

plugin/vim-python-docstring.vim

@@ -1,2 +1,3 @@
 command! -nargs=0 Docstring call vimpythondocstring#Full()
+command! -nargs=0 DocstringTypes call vimpythondocstring#FullTypes()
 command! -nargs=0 DocstringLine call vimpythondocstring#Oneline()

python/asthelper.py

@@ -84,12 +84,16 @@ def _handle_functions(self, node):
 
         if self.parent:
             for arg in node.args.args:
-                self.arguments.append(arg.arg)
+                type_hint = None
+                if arg.annotation is not None:
+                    type_hint = ast.unparse(arg.annotation)
+                self.arguments.append({'arg': arg.arg, 'type': type_hint})
             if len(self.arguments) > 0 and (self.arguments[0] == 'self' or self.arguments[0] == 'cls'):
                 self.arguments.pop(0)
 
             self.returns = new_visitor.returns
             self.yields = new_visitor.yields
+            print(self.arguments)
 
     def visit_Raise(self, node):
         r = RaiseNameCollector()

python/pydocstring.py

@@ -47,10 +47,10 @@ def _docstring_helper(self, obj_indent, docstring):
 
         return '\n'.join(lines)
 
-    def get_method_docstring(self, method_indent, args, returns, yields, raises):
+    def get_method_docstring(self, method_indent, args, returns, yields, raises, print_hints=False):
         with open(os.path.join(self.location, '..', 'styles/{}-{}.txt'.format(self.style, 'method')), 'r') as f:
             self.template = ibis.Template(f.read())
-        docstring = self.template.render(indent=self.indent, args=args,
+        docstring = self.template.render(indent=self.indent, args=args, hints=print_hints,
                                          raises=raises, returns=returns, yields=yields)
         return self._docstring_helper(method_indent, docstring)
 
@@ -71,13 +71,13 @@ class ObjectWithDocstring(abc.ABC):
 
     """
 
-    def __init__(self, env, templater, style='google'):
+    def __init__(self, env, templater):
         self.starting_line = env.current_line_nr
         self.env = env
         self.templater = templater
 
     @abc.abstractmethod
-    def write_docstring(self):
+    def write_docstring(self, *args, **kwargs):
         """ Method to create a docstring for appropriate object
 
         Writes the docstring to correct lines in `self.env` object.
@@ -88,7 +88,7 @@ def _get_sig(self):
         lines = []
         lines_it = self.env.lines_following_cursor()
         sig_line, first_line = next(lines_it)
-        indent = re.findall('^(\s*)', first_line)[0]
+        indent = re.findall(r'^(\s*)', first_line)[0]
 
         lines.append(first_line)
 
@@ -108,7 +108,7 @@ def _object_tree(self):
 
         lines.append(first_line)
 
-        obj_indent = re.findall('^(\s*)', first_line)[0]
+        obj_indent = re.findall(r'^(\s*)', first_line)[0]
         expected_indent = concat_(obj_indent, self.env.python_indent)
 
         valid_sig, _ = self._is_valid(first_line)
@@ -182,8 +182,8 @@ def write_simple_docstring(self):
 
 class MethodController(ObjectWithDocstring):
 
-    def __init__(self, env, templater, style='google'):
-        super().__init__(env, templater, style)
+    def __init__(self, env, templater):
+        super().__init__(env, templater)
 
     def _process_tree(self, tree):
         v = MethodVisitor()
@@ -193,11 +193,11 @@ def _process_tree(self, tree):
         return args, v.returns, v.yields, raises
 
     # TODO: set cursor on appropriate position to fill the docstring
-    def write_docstring(self):
+    def write_docstring(self, print_hints=False):
         sig_line, method_indent, tree = self._object_tree()
         args, returns, yields, raises = self._process_tree(tree)
         docstring = self.templater.get_method_docstring(
-            method_indent, args, returns, yields, raises)
+            method_indent, args, returns, yields, raises, print_hints)
         self.env.append_after_line(sig_line, docstring)
 
     def _arguments(self, tree):
@@ -214,8 +214,8 @@ def _arguments(self, tree):
 
 class ClassController(ObjectWithDocstring):
 
-    def __init__(self, env, templater, style='google'):
-        super().__init__(env, templater, style)
+    def __init__(self, env, templater):
+        super().__init__(env, templater)
 
     def _process_tree(self, tree):
         x = ClassInstanceNameExtractor()
@@ -225,7 +225,7 @@ def _process_tree(self, tree):
         att = list(v.attributes)
         return att
 
-    def write_docstring(self):
+    def write_docstring(self, *args, **kwargs):
         sig_line, class_indent, tree = self._object_tree()
         attr = self._process_tree(tree)
         docstring = self.templater.get_class_docstring(class_indent, attr)
@@ -240,33 +240,34 @@ def __init__(self):
         style = env.python_style
         indent = env.python_indent
         location = env.plugin_root_dir
-        templater = Templater(location, indent, style)
+        templater = Templater(location, indent=indent, style=style)
 
-        self.obj_controller = self._controller_factory(env, templater, style)
+        self.obj_controller = self._controller_factory(env, templater)
 
-    def _controller_factory(self, env, templater, style):
+    def _controller_factory(self, env, templater):
         line = env.current_line
-        first_word = re.match('^\s*(\w+).*', line).groups()[0]
+        first_word = re.match(r'^\s*(\w+).*', line).groups()[0]
         if first_word == 'def':
-            return MethodController(env, templater, style=style)
+            return MethodController(env, templater)
         elif first_word == 'class':
-            return ClassController(env, templater, style=style)
+            return ClassController(env, templater)
         elif first_word == 'async':
-            second_word_catch = re.match('^\s*\w+\s+(\w+).*', line)
+            second_word_catch = re.match(r'^\s*\w+\s+(\w+).*', line)
             if second_word_catch:
                 second_word = second_word_catch.groups()[0]
                 if second_word == 'def':
-                    return MethodController(env, templater, style=style)
+                    return MethodController(env, templater)
 
         raise DocstringUnavailable(
             'Docstring cannot be created for selected object')
 
-    def full_docstring(self):
+    def full_docstring(self, print_hints=False):
         """ Writes docstring containing arguments, returns, raises, ... """
         try:
-            self.obj_controller.write_docstring()
+            self.obj_controller.write_docstring(print_hints=print_hints)
         except Exception as e:
             print(concat_('Doctring ERROR: ', e))
+            raise e
 
     def oneline_docstring(self):
         """ Writes only a one-line empty docstring """

python/vimenv.py

@@ -97,7 +97,6 @@ def append_after_line(self, line_nr, text):
 
     def lines_following_cursor(self):
         import vim
-        lines = []
         buffer = vim.current.buffer
         cursor_row = vim.current.window.cursor[0] - 1
         current_row = cursor_row

styles/epytext-method.txt

@@ -1,6 +1,7 @@
 """
-{% if args|len > 0 %}{% for a in args %}@param {{a}}:
-{% endfor %}{% endif %}{% if returns %}@return:
+{% if args|len > 0 %}{% for a in args %}@param {{a.arg}}:
+{% endfor %}{% if hints %}{% for a in args %}{% if a.type %}@type {{a.arg}}: {{a.type}}
+{% endif %}{% endfor %}{% endif %}{% endif %}{% if returns %}@return:
 {% endif %}{% if yields %}@yield:
 {% endif %}{% if raises|len > 0 %}{% for a in raises %}@raise {{a}}:
 {% endfor %}{% endif %}"""

styles/google-method.txt

@@ -1,7 +1,7 @@
 """
 {% if args|len > 0 %}
 Args:
-{% for a in args %}{{indent}}{{a}}:
+{% for a in args %}{{indent}}{{a.arg}}{% if hints and a.type %} ({{a.type}}){% endif %}:
 {% endfor %}{% endif %}{% if returns %}
 Returns:
 {{indent}}
@@ -10,6 +10,6 @@ Yields:
 {{indent}}
 {% endif %}{% if raises|len > 0 %}
 Raises:
-{% for a in raises %}{{indent}}{{a}}:
+{% for r in raises %}{{indent}}{{r}}:
 {% endfor %}{% endif %}
 """

styles/numpy-method.txt

@@ -2,7 +2,7 @@
 {% if args|len > 0 %}
 Parameters
 ----------
-{% for a in args %}{{indent}}{{a}} :
+{% for a in args %}{{indent}}{{a.arg}} :{% if hints and a.type %} {{a.type}}{% endif %}
 {% endfor %}{% endif %}{% if returns %}
 Returns
 -------

styles/rest-method.txt

@@ -1,6 +1,7 @@
 """
-{% if args|len > 0 %}{% for a in args %}:param {{a}}:
-{% endfor %}{% endif %}{% if returns %}:returns:
+{% if args|len > 0 %}{% for a in args %}:param {{a.arg}}:
+{% endfor %}{% if hints %}{% for a in args %}{% if a.type %}:type {{a.arg}}: {{a.type}}
+{% endif %}{% endfor %}{% endif %}{% endif %}{% if returns %}:returns:
 {% endif %}{% if yields %}:yields:
 {% endif %}{% if raises|len > 0 %}{% for a in raises %}:raises {{a}}:
 {% endfor %}{% endif %}"""