pycode: Support type annotations for variables

sphinx/pycode/__init__.py

@@ -142,9 +142,10 @@ class ModuleAnalyzer:
             self.code = source.read()
 
         # will be filled by parse()
-        self.attr_docs = None   # type: Dict[Tuple[str, str], List[str]]
-        self.tagorder = None    # type: Dict[str, int]
-        self.tags = None        # type: Dict[str, Tuple[str, int, int]]
+        self.annotations = None  # type: Dict[Tuple[str, str], str]
+        self.attr_docs = None    # type: Dict[Tuple[str, str], List[str]]
+        self.tagorder = None     # type: Dict[str, int]
+        self.tags = None         # type: Dict[str, Tuple[str, int, int]]
 
     def parse(self) -> None:
         """Parse the source code."""
@@ -159,6 +160,7 @@ class ModuleAnalyzer:
                 else:
                     self.attr_docs[scope] = ['']
 
+            self.annotations = parser.annotations
             self.tags = parser.definitions
             self.tagorder = parser.deforders
         except Exception as exc:

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
 

tests/test_pycode_parser.py

@@ -105,6 +105,9 @@ def test_annotated_assignment_py36():
     assert parser.comments == {('', 'a'): 'comment',
                                ('', 'b'): 'string on next line',
                                ('', 'c'): 'comment'}
+    assert parser.annotations == {('', 'a'): 'str',
+                                  ('', 'b'): 'int',
+                                  ('', 'c'): 'int'}
     assert parser.definitions == {}