get from-import dependencies

2 files changed, 31 insertions, 11 deletions

diff --git a/pyreverse/diadefslib.py b/pyreverse/diadefslib.py
index 9511a13..6d951f4 100644
--- a/pyreverse/diadefslib.py
+++ b/pyreverse/diadefslib.py
@@ -210,6 +210,12 @@ class DefaultDiadefGenerator(LocalsVisitor, OptionHandler):
         title = self.get_title(node)
         self.classdiagram.add_object(title, node)
 
+    def visit_from(self, node):
+        """visit astng.From  and catch modules for package diagram
+        """
+        if self.pkgdiagram:
+            self.pkgdiagram.add_depend_relation( node, node.modname )
+
 
 class ClassDiadefGenerator(OptionHandler):
     """generate a class diagram definition including all classes related to a
@@ -258,7 +264,6 @@ class ClassDiadefGenerator(OptionHandler):
         # association
         for name, ass_nodes in klass_node.instance_attrs_type.items():
             for ass_node in ass_nodes:
-                # XXX could find here class attributes and their type
                 if isinstance(ass_node, astng.Instance):
                     ass_node = ass_node._proxied
                 if not isinstance(ass_node, astng.Class) \
diff --git a/pyreverse/diagrams.py b/pyreverse/diagrams.py
index 1f7e653..5f2c3db 100644
--- a/pyreverse/diagrams.py
+++ b/pyreverse/diagrams.py
@@ -58,7 +58,8 @@ class ClassDiagram(Figure, FilterMixIn):
         self.objects = []
         self.relationships = {}
         self._nodes = {}
-        
+        self.depends = []
+
     def add_relationship(self, from_object, to_object, relation_type, name=None):
         """create a relation ship
         """
@@ -107,7 +108,8 @@ class ClassDiagram(Figure, FilterMixIn):
             if isinstance(ass_node, astng.Class) \
                 and hasattr(ass_node, "name") and not self.has_node(ass_node):
                 if ass_node.name not in names:
-                    names.append(ass_node.name)
+                    ass_name = ass_node.name
+                    names.append(ass_name)
         return names
 
     def nodes(self):
@@ -142,8 +144,8 @@ class ClassDiagram(Figure, FilterMixIn):
         """
         for obj in self.classes():
             node = obj.node
-            obj.attrs = self.get_attrs(obj.node)
-            obj.methods = self.get_methods(obj.node)
+            obj.attrs = self.get_attrs(node)
+            obj.methods = self.get_methods(node)
             # shape
             if is_interface(node):
                 obj.shape = 'interface'
@@ -190,25 +192,38 @@ class PackageDiagram(ClassDiagram):
             if mod.node.name == name:
                 return mod
         raise KeyError(name)
-    
+
+    def add_depend_relation(self, node, from_module):
+        """add dependencies created by from-imports
+        """
+        mod_name = node.root().name
+        obj = self.module( mod_name )
+        if from_module not in obj.node.depends:
+            obj.node.depends.append(from_module)
+
     def extract_relationships(self):
         """extract relation ships between nodes in the diagram
         """
         ClassDiagram.extract_relationships(self)
         for obj in self.classes():
-            node = obj.node
             # ownership
             try:
-                mod = self.object_from_node(node.root())
+                mod = self.object_from_node(obj.node.root())
                 self.add_relationship(obj, mod, 'ownership')
             except KeyError:
                 continue
         for obj in self.modules():
             obj.shape = 'package'
             # dependencies
-            for dep in obj.node.depends:
+            for dep_name in obj.node.depends:
                 try:
-                    dep = self.module(dep)
+                    dep = self.module(dep_name)
                 except KeyError:
-                    continue
+                    # relative imports
+                    package = obj.node.root().name.rsplit('.', 1)[0]
+                    dep_name = "%s.%s" % (package, dep_name)
+                    try:
+                        dep = self.module(dep_name)
+                    except KeyError:
+                        continue
                 self.add_relationship(obj, dep, 'depends')