pyreverse: take into account class attributes which are stored in locals_type

2 files changed, 8 insertions, 4 deletions

diff --git a/pyreverse/diadefslib.py b/pyreverse/diadefslib.py
index 0496ac2..a8b15d3 100644
--- a/pyreverse/diadefslib.py
+++ b/pyreverse/diadefslib.py
@@ -95,7 +95,8 @@ class DiaDefGenerator:
         """return associated nodes of a class node"""
         if level == 0:
             return
-        for ass_nodes in klass_node.instance_attrs_type.values():
+        for ass_nodes in klass_node.instance_attrs_type.values() + \
+                         klass_node.locals_type.values():
             for ass_node in ass_nodes:
                 if isinstance(ass_node, astng.Instance):
                     ass_node = ass_node._proxied
diff --git a/pyreverse/diagrams.py b/pyreverse/diagrams.py
index 87afa4c..5725fd7 100644
--- a/pyreverse/diagrams.py
+++ b/pyreverse/diagrams.py
@@ -80,7 +80,8 @@ class ClassDiagram(Figure, FilterMixIn):
     def get_attrs(self, node):
         """return visible attributes, possibly with class name"""
         attrs = []
-        for node_name, ass_nodes in node.instance_attrs_type.items():
+        for node_name, ass_nodes in node.instance_attrs_type.items() + \
+                                node.locals_type.items():
             if not self.show_attr(node_name):
                 continue
             names = self.class_names(ass_nodes)
@@ -169,7 +170,8 @@ class ClassDiagram(Figure, FilterMixIn):
                 except KeyError:
                     continue
             # associations link
-            for name, values in node.instance_attrs_type.items():
+            for name, values in node.instance_attrs_type.items() + \
+                                node.locals_type.items():
                 for value in values:
                     if value is astng.YES:
                         continue
@@ -180,7 +182,8 @@ class ClassDiagram(Figure, FilterMixIn):
                         self.add_relationship(ass_obj, obj, 'association', name)
                     except KeyError:
                         continue
-        
+
+
 class PackageDiagram(ClassDiagram):
     """package diagram handling
     """