issue #217: ansible: working module_finder.py

ansible_mitogen/module_finder.py

@@ -1,49 +1,30 @@
 
+from __future__ import absolute_import
 import imp
 import os
 import sys
-import mitogen.master
-
-
-class Name(object):
-    def __str__(self):
-        return self.identifier
-
-    def __repr__(self):
-        return 'Name(%r)' % (self.identifier,)
 
-    def __init__(self, identifier):
-        self.identifier = identifier
-
-    def head(self):
-        head, _, tail = self.identifier.partition('.')
-        return head
-
-    def tail(self):
-        head, _, tail = self.identifier.partition('.')
-        return tail
+import mitogen.master
+import ansible.module_utils
 
-    def pop_n(self, level):
-        name = self.identifier
-        for _ in xrange(level):
-            if '.' not in name:
-                return None
-            name, _, _ = self.identifier.rpartition('.')
-        return Name(name)
 
-    def append(self, part):
-        return Name('%s.%s' % (self.identifier, part))
+PREFIX = 'ansible.module_utils.'
 
 
 class Module(object):
     def __init__(self, name, path, kind=imp.PY_SOURCE, parent=None):
-        self.name = Name(name)
+        self.name = name
         self.path = path
-        if kind == imp.PKG_DIRECTORY:
-            self.path = os.path.join(self.path, '__init__.py')
         self.kind = kind
+        self.is_pkg = kind == imp.PKG_DIRECTORY
         self.parent = parent
 
+    def __hash__(self):
+        return hash(self.path)
+
+    def __eq__(self, other):
+        return self.path == other.path
+
     def fullname(self):
         bits = [str(self.name)]
         while self.parent:
@@ -51,16 +32,6 @@ class Module(object):
             self = self.parent
         return '.'.join(reversed(bits))
 
-    def __repr__(self):
-        return 'Module(%r, path=%r, parent=%r)' % (
-            self.name,
-            self.path,
-            self.parent,
-        )
-
-    def dirname(self):
-        return os.path.dirname(self.path)
-
     def code(self):
         fp = open(self.path)
         try:
@@ -73,8 +44,9 @@ def find(name, path=(), parent=None):
     """
     (Name, search path) -> Module instance or None.
     """
+    head, _, tail = name.partition('.')
     try:
-        tup = imp.find_module(name.head(), list(path))
+        tup = imp.find_module(head, list(path))
     except ImportError:
         return parent
 
@@ -82,56 +54,53 @@ def find(name, path=(), parent=None):
     if fp:
         fp.close()
 
-    module = Module(name.head(), path, kind, parent)
-    if name.tail():
-        return find_relative(module, Name(name.tail()), path)
+    if kind == imp.PKG_DIRECTORY:
+        path = os.path.join(path, '__init__.py')
+    module = Module(head, path, kind, parent)
+    if tail:
+        return find_relative(module, tail, path)
     return module
 
 
 def find_relative(parent, name, path=()):
-    path = [parent.dirname()] + list(path)
+    path = [os.path.dirname(parent.path)] + list(path)
     return find(name, path, parent=parent)
 
 
-def path_pop(s, n):
-    return os.pathsep.join(s.split(os.pathsep)[-n:])
-
-
-def scan(module, path):
-    scanner = mitogen.master.scan_code_imports(module.code())
-    for level, modname_s, fromlist in scanner:
-        modname = Name(modname_s)
-        if level == -1:
-            imported = find_relative(module, modname, path)
-        elif level:
-            subpath = [path_pop(module.dirname(), level)] + list(path)
-            imported = find(modname.pop_n(level), subpath)
-        else:
-            imported = find(modname.pop_n(level), path)
-
-        if imported and mitogen.master.is_stdlib_path(imported.path):
-            continue
-
-        if imported and fromlist:
-            have = False
-            for fromname_s in fromlist:
-                fromname = modname.append(fromname_s)
-                f_imported = find_relative(imported, fromname, path)
-                if f_imported and f_imported.fullname() == fromname.identifier:
-                    have = True
-                    yield fromname, f_imported, None
-            if have:
-                continue
+def scan_fromlist(code):
+    for level, modname_s, fromlist in mitogen.master.scan_code_imports(code):
+        for name in fromlist:
+            yield level, '%s.%s' % (modname_s, name)
+        if not fromlist:
+            yield level, modname_s
 
-        if imported:
-            yield modname, imported
 
+def scan(module_name, module_path, search_path):
+    module = Module(module_name, module_path)
+    stack = [module]
+    seen = set()
 
-module = Module(name='ansible_module_apt', path='/Users/dmw/src/mitogen/.venv/lib/python2.7/site-packages/ansible/modules/packaging/os/apt.py')
-path = tuple(sys.path)
-path = ('/Users/dmw/src/ansible/lib',) + path
+    while stack:
+        module = stack.pop(0)
+        for level, fromname in scan_fromlist(module.code()):
+            if not fromname.startswith(PREFIX):
+                continue
 
+            imported = find(fromname[len(PREFIX):], search_path)
+            if imported is None or imported in seen:
+                continue
 
-from pprint import pprint
-for name, imported in scan(module, sys.path):
-    print '%s: %s' % (name, imported and (str(name) == imported.fullname()))
+            seen.add(imported)
+            parent = imported.parent
+            while parent:
+                module = Module(name=parent.fullname(), path=parent.path,
+                                kind=parent.kind)
+                if module not in seen:
+                    seen.add(module)
+                    stack.append(module)
+                parent = parent.parent
+
+    return sorted(
+        (PREFIX + module.fullname(), module.path, module.is_pkg)
+        for module in seen
+    )