diff options
-rw-r--r-- | pyreverse/diadefslib.py | 5 | ||||
-rw-r--r-- | pyreverse/diagrams.py | 28 |
2 files changed, 22 insertions, 11 deletions
diff --git a/pyreverse/diadefslib.py b/pyreverse/diadefslib.py index b36acdb..8a98d68 100644 --- a/pyreverse/diadefslib.py +++ b/pyreverse/diadefslib.py @@ -198,7 +198,7 @@ class DiaDefGenerator: """return associated nodes of a class node""" if level == 0: return - for name, ass_nodes in klass_node.instance_attrs_type.items(): + for ass_nodes in klass_node.instance_attrs_type.values(): for ass_node in ass_nodes: if isinstance(ass_node, astng.Instance): ass_node = ass_node._proxied @@ -274,7 +274,8 @@ class DefaultDiadefGenerator(LocalsVisitor, DiaDefGenerator): """visit astng.From and catch modules for package diagram """ if self.pkgdiagram: - self.pkgdiagram.add_depend_relation( node, node.modname ) + self.pkgdiagram.add_from_depend(node, node.modname) + class ClassDiadefGenerator(DiaDefGenerator): """generate a class diagram definition including all classes related to a diff --git a/pyreverse/diagrams.py b/pyreverse/diagrams.py index 639ed17..426d4e7 100644 --- a/pyreverse/diagrams.py +++ b/pyreverse/diagrams.py @@ -198,7 +198,23 @@ class PackageDiagram(ClassDiagram): return mod raise KeyError(name) - def add_depend_relation(self, node, from_module): + def get_module(self, name, node): + """return a module by its name, looking also for relative imports; + raise KeyError if not found + """ + for mod in self.modules(): + mod_name = mod.node.name + if mod_name == name: + return mod + #search for fullname of relative import modules + package = node.root().name + if mod_name == "%s.%s" % (package, name): + return mod + if mod_name == "%s.%s" % (package.rsplit('.', 1)[0], name): + return mod + raise KeyError(name) + + def add_from_depend(self, node, from_module): """add dependencies created by from-imports """ mod_name = node.root().name @@ -222,13 +238,7 @@ class PackageDiagram(ClassDiagram): # dependencies for dep_name in obj.node.depends: try: - dep = self.module(dep_name) + dep = self.get_module(dep_name, obj.node) except KeyError: - # 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 + continue self.add_relationship(obj, dep, 'depends') |