Add type annotations to pyreverse dot files

- Add a try/except to deal with a possible InferenceError when using
NodeNG.infer method
- Create a function in utils and so remove repeated logic in inspector.py
- Add unittests to check the InferenceError logic
- Adjust the types in function input

Author: mbyrnepr2

pylint/pyreverse/inspector.py

@@ -205,38 +205,30 @@ def visit_assignname(self, node):
             # the name has been defined as 'global' in the frame and belongs
             # there.
             frame = node.root()
-        try:
-            if not hasattr(frame, "locals_type"):
-                # If the frame doesn't have a locals_type yet,
-                # it means it wasn't yet visited. Visit it now
-                # to add what's missing from it.
-                if isinstance(frame, astroid.ClassDef):
-                    self.visit_classdef(frame)
-                elif isinstance(frame, astroid.FunctionDef):
-                    self.visit_functiondef(frame)
-                else:
-                    self.visit_module(frame)
-
-            current = frame.locals_type[node.name]
-            ann = utils.get_annotation(node)
-            values = {ann} if ann else set(node.infer())
-            frame.locals_type[node.name] = list(set(current) | values)
-        except astroid.InferenceError:
-            pass
+        if not hasattr(frame, "locals_type"):
+            # If the frame doesn't have a locals_type yet,
+            # it means it wasn't yet visited. Visit it now
+            # to add what's missing from it.
+            if isinstance(frame, astroid.ClassDef):
+                self.visit_classdef(frame)
+            elif isinstance(frame, astroid.FunctionDef):
+                self.visit_functiondef(frame)
+            else:
+                self.visit_module(frame)
+
+        current = frame.locals_type[node.name]
+        frame.locals_type[node.name] = list(set(current) | utils.infer_node(node))
 
     @staticmethod
     def handle_assignattr_type(node, parent):
         """handle an astroid.assignattr node
 
         handle instance_attrs_type
         """
-        ann = utils.get_annotation(node)
-        try:
-            values = {ann} if ann else set(node.infer())
-            current = set(parent.instance_attrs_type[node.attrname])
-            parent.instance_attrs_type[node.attrname] = list(current | values)
-        except astroid.InferenceError:
-            pass
+        current = set(parent.instance_attrs_type[node.attrname])
+        parent.instance_attrs_type[node.attrname] = list(
+            current | utils.infer_node(node)
+        )
 
     def visit_import(self, node):
         """visit an astroid.Import node

pylint/pyreverse/utils.py

@@ -228,7 +228,7 @@ def get_annotation_label(ann: Union[astroid.Name, astroid.Subscript]) -> str:
 
 
 def get_annotation(
-    node: Union[astroid.AnnAssign, astroid.AssignAttr]
+    node: Union[astroid.AssignAttr, astroid.AssignName]
 ) -> Optional[Union[astroid.Name, astroid.Subscript]]:
     """return the annotation for `node`"""
     ann = None
@@ -244,7 +244,11 @@ def get_annotation(
     else:
         return ann
 
-    default, *_ = node.infer()
+    try:
+        default, *_ = node.infer()
+    except astroid.InferenceError:
+        default = ""
+
     label = get_annotation_label(ann)
     if ann:
         label = (
@@ -256,3 +260,16 @@ def get_annotation(
     if label:
         ann.name = label
     return ann
+
+
+def infer_node(node: Union[astroid.AssignAttr, astroid.AssignName]) -> set:
+    """Return a set containing the node annotation if it exists
+    otherwise return a set of the inferred types using the NodeNG.infer method"""
+
+    ann = get_annotation(node)
+    if ann:
+        return {ann}
+    try:
+        return set(node.infer())
+    except astroid.InferenceError:
+        return set()

tests/unittest_pyreverse_writer.py

@@ -22,13 +22,14 @@
 import codecs
 import os
 from difflib import unified_diff
+from unittest.mock import patch
 
 import astroid
 import pytest
 
 from pylint.pyreverse.diadefslib import DefaultDiadefGenerator, DiadefsHandler
 from pylint.pyreverse.inspector import Linker, project_from_files
-from pylint.pyreverse.utils import get_annotation, get_visibility
+from pylint.pyreverse.utils import get_annotation, get_visibility, infer_node
 from pylint.pyreverse.writer import DotWriter
 
 _DEFAULTS = {
@@ -176,3 +177,29 @@ class A:
             got = get_annotation(assign_attr).name
             assert isinstance(assign_attr, astroid.AssignAttr)
             assert got == label, f"got {got} instead of {label} for value {node}"
+
+
+@patch("pylint.pyreverse.utils.get_annotation")
+@patch("astroid.node_classes.NodeNG.infer", side_effect=astroid.InferenceError)
+def test_infer_node_1(mock_infer, mock_get_annotation):
+    """Return set() when astroid.InferenceError is raised and an annotation has
+    not been returned
+    """
+    mock_get_annotation.return_value = None
+    node = astroid.extract_node("a: str = 'mystr'")
+    mock_infer.return_value = "x"
+    assert infer_node(node) == set()
+    assert mock_infer.called
+
+
+@patch("pylint.pyreverse.utils.get_annotation")
+@patch("astroid.node_classes.NodeNG.infer")
+def test_infer_node_2(mock_infer, mock_get_annotation):
+    """Return set(node.infer()) when InferenceError is not raised and an
+    annotation has not been returned
+    """
+    mock_get_annotation.return_value = None
+    node = astroid.extract_node("a: str = 'mystr'")
+    mock_infer.return_value = "x"
+    assert infer_node(node) == set("x")
+    assert mock_infer.called