pycode: Support type annotations for variables

1 files changed, 25 insertions, 2 deletions

diff --git a/sphinx/pycode/parser.py b/sphinx/pycode/parser.py
index 855167ba7..4ec1a844d 100644
--- a/sphinx/pycode/parser.py
+++ b/sphinx/pycode/parser.py
@@ -7,7 +7,6 @@
     :copyright: Copyright 2007-2020 by the Sphinx team, see AUTHORS.
     :license: BSD, see LICENSE for details.
 """
-import ast
 import inspect
 import itertools
 import re
@@ -17,6 +16,9 @@ from token import NAME, NEWLINE, INDENT, DEDENT, NUMBER, OP, STRING
 from tokenize import COMMENT, NL
 from typing import Any, Dict, List, Tuple
 
+from sphinx.pycode.ast import ast  # for py37 or older
+from sphinx.pycode.ast import parse, unparse
+
 
 comment_re = re.compile('^\\s*#: ?(.*)\r?\n?$')
 indent_re = re.compile('^\\s*$')
@@ -226,6 +228,7 @@ class VariableCommentPicker(ast.NodeVisitor):
         self.current_classes = []       # type: List[str]
         self.current_function = None    # type: ast.FunctionDef
         self.comments = {}              # type: Dict[Tuple[str, str], str]
+        self.annotations = {}           # type: Dict[Tuple[str, str], str]
         self.previous = None            # type: ast.AST
         self.deforders = {}             # type: Dict[str, int]
         super().__init__()
@@ -254,6 +257,18 @@ class VariableCommentPicker(ast.NodeVisitor):
 
         self.comments[(context, name)] = comment
 
+    def add_variable_annotation(self, name: str, annotation: ast.AST) -> None:
+        if self.current_function:
+            if self.current_classes and self.context[-1] == "__init__":
+                # store variable comments inside __init__ method of classes
+                context = ".".join(self.context[:-1])
+            else:
+                return
+        else:
+            context = ".".join(self.context)
+
+        self.annotations[(context, name)] = unparse(annotation)
+
     def get_self(self) -> ast.arg:
         """Returns the name of first argument if in function."""
         if self.current_function and self.current_function.args.args:
@@ -295,6 +310,12 @@ class VariableCommentPicker(ast.NodeVisitor):
         except TypeError:
             return  # this assignment is not new definition!
 
+        # record annotation
+        annotation = getattr(node, 'annotation', None)
+        if annotation:
+            for varname in varnames:
+                self.add_variable_annotation(varname, annotation)
+
         # check comments after assignment
         parser = AfterCommentParser([current_line[node.col_offset:]] +
                                     self.buffers[node.lineno:])
@@ -468,6 +489,7 @@ class Parser:
     def __init__(self, code: str, encoding: str = 'utf-8') -> None:
         self.code = filter_whitespace(code)
         self.encoding = encoding
+        self.annotations = {}       # type: Dict[Tuple[str, str], str]
         self.comments = {}          # type: Dict[Tuple[str, str], str]
         self.deforders = {}         # type: Dict[str, int]
         self.definitions = {}       # type: Dict[str, Tuple[str, int, int]]
@@ -479,9 +501,10 @@ class Parser:
 
     def parse_comments(self) -> None:
         """Parse the code and pick up comments."""
-        tree = ast.parse(self.code)
+        tree = parse(self.code)
         picker = VariableCommentPicker(self.code.splitlines(True), self.encoding)
         picker.visit(tree)
+        self.annotations = picker.annotations
         self.comments = picker.comments
         self.deforders = picker.deforders