Tools/freeze/modulefinder.py

   1 """Find modules used by a script, using introspection."""
   2
   3 import dis
   4 import imp
   5 import marshal
   6 import os
   7 import re
   8 import string
   9 import sys
  10
  11 if sys.platform=="win32":
  12     # On Windows, we can locate modules in the registry with
  13     # the help of the win32api package.
  14     try:
  15         import win32api
  16     except ImportError:
  17         print "The win32api module is not available - modules listed"
  18         print "in the registry will not be found."
  19         win32api = None
  20
  21
  22 IMPORT_NAME = dis.opname.index('IMPORT_NAME')
  23 IMPORT_FROM = dis.opname.index('IMPORT_FROM')
  24
  25 # Modulefinder does a good job at simulating Python's, but it can not
  26 # handle __path__ modifications packages make at runtime.  Therefore there
  27 # is a mechanism whereby you can register extra paths in this map for a
  28 # package, and it will be honoured.
  29
  30 # Note this is a mapping is lists of paths.
  31 packagePathMap = {}
  32
  33 # A Public interface
  34 def AddPackagePath(packagename, path):
  35     paths = packagePathMap.get(packagename, [])
  36     paths.append(path)
  37     packagePathMap[packagename] = paths
  38
  39 class Module:
  40
  41     def __init__(self, name, file=None, path=None):
  42         self.__name__ = name
  43         self.__file__ = file
  44         self.__path__ = path
  45         self.__code__ = None
  46
  47     def __repr__(self):
  48         s = "Module(%s" % `self.__name__`
  49         if self.__file__ is not None:
  50             s = s + ", %s" % `self.__file__`
  51         if self.__path__ is not None:
  52             s = s + ", %s" % `self.__path__`
  53         s = s + ")"
  54         return s
  55
  56
  57 class ModuleFinder:
  58
  59     def __init__(self, path=None, debug=0, excludes = []):
  60         if path is None:
  61             path = sys.path
  62         self.path = path
  63         self.modules = {}
  64         self.badmodules = {}
  65         self.debug = debug
  66         self.indent = 0
  67         self.excludes = excludes
  68
  69     def msg(self, level, str, *args):
  70         if level <= self.debug:
  71             for i in range(self.indent):
  72                 print "   ",
  73             print str,
  74             for arg in args:
  75                 print repr(arg),
  76             print
  77
  78     def msgin(self, *args):
  79         level = args[0]
  80         if level <= self.debug:
  81             self.indent = self.indent + 1
  82             apply(self.msg, args)
  83
  84     def msgout(self, *args):
  85         level = args[0]
  86         if level <= self.debug:
  87             self.indent = self.indent - 1
  88             apply(self.msg, args)
  89
  90     def run_script(self, pathname):
  91         self.msg(2, "run_script", pathname)
  92         fp = open(pathname)
  93         stuff = ("", "r", imp.PY_SOURCE)
  94         self.load_module('__main__', fp, pathname, stuff)
  95
  96     def load_file(self, pathname):
  97         dir, name = os.path.split(pathname)
  98         name, ext = os.path.splitext(name)
  99         fp = open(pathname)
 100         stuff = (ext, "r", imp.PY_SOURCE)
 101         self.load_module(name, fp, pathname, stuff)
 102
 103     def import_hook(self, name, caller=None, fromlist=None):
 104         self.msg(3, "import_hook", name, caller, fromlist)
 105         parent = self.determine_parent(caller)
 106         q, tail = self.find_head_package(parent, name)
 107         m = self.load_tail(q, tail)
 108         if not fromlist:
 109             return q
 110         if m.__path__:
 111             self.ensure_fromlist(m, fromlist)
 112
 113     def determine_parent(self, caller):
 114         self.msgin(4, "determine_parent", caller)
 115         if not caller:
 116             self.msgout(4, "determine_parent -> None")
 117             return None
 118         pname = caller.__name__
 119         if caller.__path__:
 120             parent = self.modules[pname]
 121             assert caller is parent
 122             self.msgout(4, "determine_parent ->", parent)
 123             return parent
 124         if '.' in pname:
 125             i = string.rfind(pname, '.')
 126             pname = pname[:i]
 127             parent = self.modules[pname]
 128             assert parent.__name__ == pname
 129             self.msgout(4, "determine_parent ->", parent)
 130             return parent
 131         self.msgout(4, "determine_parent -> None")
 132         return None
 133
 134     def find_head_package(self, parent, name):
 135         self.msgin(4, "find_head_package", parent, name)
 136         if '.' in name:
 137             i = string.find(name, '.')
 138             head = name[:i]
 139             tail = name[i+1:]
 140         else:
 141             head = name
 142             tail = ""
 143         if parent:
 144             qname = "%s.%s" % (parent.__name__, head)
 145         else:
 146             qname = head
 147         q = self.import_module(head, qname, parent)
 148         if q:
 149             self.msgout(4, "find_head_package ->", (q, tail))
 150             return q, tail
 151         if parent:
 152             qname = head
 153             parent = None
 154             q = self.import_module(head, qname, parent)
 155             if q:
 156                 self.msgout(4, "find_head_package ->", (q, tail))
 157                 return q, tail
 158         self.msgout(4, "raise ImportError: No module named", qname)
 159         raise ImportError, "No module named " + qname
 160
 161     def load_tail(self, q, tail):
 162         self.msgin(4, "load_tail", q, tail)
 163         m = q
 164         while tail:
 165             i = string.find(tail, '.')
 166             if i < 0: i = len(tail)
 167             head, tail = tail[:i], tail[i+1:]
 168             mname = "%s.%s" % (m.__name__, head)
 169             m = self.import_module(head, mname, m)
 170             if not m:
 171                 self.msgout(4, "raise ImportError: No module named", mname)
 172                 raise ImportError, "No module named " + mname
 173         self.msgout(4, "load_tail ->", m)
 174         return m
 175
 176     def ensure_fromlist(self, m, fromlist, recursive=0):
 177         self.msg(4, "ensure_fromlist", m, fromlist, recursive)
 178         for sub in fromlist:
 179             if sub == "*":
 180                 if not recursive:
 181                     all = self.find_all_submodules(m)
 182                     if all:
 183                         self.ensure_fromlist(m, all, 1)
 184             elif not hasattr(m, sub):
 185                 subname = "%s.%s" % (m.__name__, sub)
 186                 submod = self.import_module(sub, subname, m)
 187                 if not submod:
 188                     raise ImportError, "No module named " + subname
 189
 190     def find_all_submodules(self, m):
 191         if not m.__path__:
 192             return
 193         modules = {}
 194         suffixes = [".py", ".pyc", ".pyo"]
 195         for dir in m.__path__:
 196             try:
 197                 names = os.listdir(dir)
 198             except os.error:
 199                 self.msg(2, "can't list directory", dir)
 200                 continue
 201             for name in names:
 202                 mod = None
 203                 for suff in suffixes:
 204                     n = len(suff)
 205                     if name[-n:] == suff:
 206                         mod = name[:-n]
 207                         break
 208                 if mod and mod != "__init__":
 209                     modules[mod] = mod
 210         return modules.keys()
 211
 212     def import_module(self, partname, fqname, parent):
 213         self.msgin(3, "import_module", partname, fqname, parent)
 214         try:
 215             m = self.modules[fqname]
 216         except KeyError:
 217             pass
 218         else:
 219             self.msgout(3, "import_module ->", m)
 220             return m
 221         if self.badmodules.has_key(fqname):
 222             self.msgout(3, "import_module -> None")
 223             if parent:
 224                 self.badmodules[fqname][parent.__name__] = None
 225             return None
 226         try:
 227             fp, pathname, stuff = self.find_module(partname,
 228                                                    parent and parent.__path__)
 229         except ImportError:
 230             self.msgout(3, "import_module ->", None)
 231             return None
 232         try:
 233             m = self.load_module(fqname, fp, pathname, stuff)
 234         finally:
 235             if fp: fp.close()
 236         if parent:
 237             setattr(parent, partname, m)
 238         self.msgout(3, "import_module ->", m)
 239         return m
 240
 241     def load_module(self, fqname, fp, pathname, (suffix, mode, type)):
 242         self.msgin(2, "load_module", fqname, fp and "fp", pathname)
 243         if type == imp.PKG_DIRECTORY:
 244             m = self.load_package(fqname, pathname)
 245             self.msgout(2, "load_module ->", m)
 246             return m
 247         if type == imp.PY_SOURCE:
 248             co = compile(fp.read()+'\n', pathname, 'exec')
 249         elif type == imp.PY_COMPILED:
 250             if fp.read(4) != imp.get_magic():
 251                 self.msgout(2, "raise ImportError: Bad magic number", pathname)
 252                 raise ImportError, "Bad magic number in %s", pathname
 253             fp.read(4)
 254             co = marshal.load(fp)
 255         else:
 256             co = None
 257         m = self.add_module(fqname)
 258         m.__file__ = pathname
 259         if co:
 260             m.__code__ = co
 261             self.scan_code(co, m)
 262         self.msgout(2, "load_module ->", m)
 263         return m
 264
 265     def scan_code(self, co, m):
 266         code = co.co_code
 267         n = len(code)
 268         i = 0
 269         lastname = None
 270         while i < n:
 271             c = code[i]
 272             i = i+1
 273             op = ord(c)
 274             if op >= dis.HAVE_ARGUMENT:
 275                 oparg = ord(code[i]) + ord(code[i+1])*256
 276                 i = i+2
 277             if op == IMPORT_NAME:
 278                 name = lastname = co.co_names[oparg]
 279                 if not self.badmodules.has_key(lastname):
 280                     try:
 281                         self.import_hook(name, m)
 282                     except ImportError, msg:
 283                         self.msg(2, "ImportError:", str(msg))
 284                         if not self.badmodules.has_key(name):
 285                             self.badmodules[name] = {}
 286                         self.badmodules[name][m.__name__] = None
 287             elif op == IMPORT_FROM:
 288                 name = co.co_names[oparg]
 289                 assert lastname is not None
 290                 if not self.badmodules.has_key(lastname):
 291                     try:
 292                         self.import_hook(lastname, m, [name])
 293                     except ImportError, msg:
 294                         self.msg(2, "ImportError:", str(msg))
 295                         fullname = lastname + "." + name
 296                         if not self.badmodules.has_key(fullname):
 297                             self.badmodules[fullname] = {}
 298                         self.badmodules[fullname][m.__name__] = None
 299             else:
 300                 lastname = None
 301         for c in co.co_consts:
 302             if isinstance(c, type(co)):
 303                 self.scan_code(c, m)
 304
 305     def load_package(self, fqname, pathname):
 306         self.msgin(2, "load_package", fqname, pathname)
 307         m = self.add_module(fqname)
 308         m.__file__ = pathname
 309         m.__path__ = [pathname]
 310
 311         # As per comment at top of file, simulate runtime __path__ additions.
 312         m.__path__ = m.__path__ + packagePathMap.get(fqname, [])
 313
 314         fp, buf, stuff = self.find_module("__init__", m.__path__)
 315         self.load_module(fqname, fp, buf, stuff)
 316         self.msgout(2, "load_package ->", m)
 317         return m
 318
 319     def add_module(self, fqname):
 320         if self.modules.has_key(fqname):
 321             return self.modules[fqname]
 322         self.modules[fqname] = m = Module(fqname)
 323         return m
 324
 325     def find_module(self, name, path):
 326         if name in self.excludes:
 327             self.msgout(3, "find_module -> Excluded")
 328             raise ImportError, name
 329
 330         if path is None:
 331             if name in sys.builtin_module_names:
 332                 return (None, None, ("", "", imp.C_BUILTIN))
 333
 334             # Emulate the Registered Module support on Windows.
 335             if sys.platform=="win32" and win32api is not None:
 336                 HKEY_LOCAL_MACHINE = 0x80000002
 337                 try:
 338                     pathname = win32api.RegQueryValue(HKEY_LOCAL_MACHINE, "Software\\Python\\PythonCore\\%s\\Modules\\%s" % (sys.winver, name))
 339                     fp = open(pathname, "rb")
 340                     # XXX - To do - remove the hard code of C_EXTENSION.
 341                     stuff = "", "rb", imp.C_EXTENSION
 342                     return fp, pathname, stuff
 343                 except win32api.error:
 344                     pass
 345
 346             path = self.path
 347         return imp.find_module(name, path)
 348
 349     def report(self):
 350         print
 351         print "  %-25s %s" % ("Name", "File")
 352         print "  %-25s %s" % ("----", "----")
 353         # Print modules found
 354         keys = self.modules.keys()
 355         keys.sort()
 356         for key in keys:
 357             m = self.modules[key]
 358             if m.__path__:
 359                 print "P",
 360             else:
 361                 print "m",
 362             print "%-25s" % key, m.__file__ or ""
 363
 364         # Print missing modules
 365         keys = self.badmodules.keys()
 366         keys.sort()
 367         for key in keys:
 368             # ... but not if they were explicitely excluded.
 369             if key not in self.excludes:
 370                 mods = self.badmodules[key].keys()
 371                 mods.sort()
 372                 print "?", key, "from", string.join(mods, ', ')
 373
 374
 375 def test():
 376     # Parse command line
 377     import getopt
 378     try:
 379         opts, args = getopt.getopt(sys.argv[1:], "dmp:qx:")
 380     except getopt.error, msg:
 381         print msg
 382         return
 383
 384     # Process options
 385     debug = 1
 386     domods = 0
 387     addpath = []
 388     exclude = []
 389     for o, a in opts:
 390         if o == '-d':
 391             debug = debug + 1
 392         if o == '-m':
 393             domods = 1
 394         if o == '-p':
 395             addpath = addpath + string.split(a, os.pathsep)
 396         if o == '-q':
 397             debug = 0
 398         if o == '-x':
 399             exclude.append(a)
 400
 401     # Provide default arguments
 402     if not args:
 403         script = "hello.py"
 404     else:
 405         script = args[0]
 406
 407     # Set the path based on sys.path and the script directory
 408     path = sys.path[:]
 409     path[0] = os.path.dirname(script)
 410     path = addpath + path
 411     if debug > 1:
 412         print "path:"
 413         for item in path:
 414             print "   ", `item`
 415
 416     # Create the module finder and turn its crank
 417     mf = ModuleFinder(path, debug, exclude)
 418     for arg in args[1:]:
 419         if arg == '-m':
 420             domods = 1
 421             continue
 422         if domods:
 423             if arg[-2:] == '.*':
 424                 mf.import_hook(arg[:-2], None, ["*"])
 425             else:
 426                 mf.import_hook(arg)
 427         else:
 428             mf.load_file(arg)
 429     mf.run_script(script)
 430     mf.report()
 431
 432
 433 if __name__ == '__main__':
 434     try:
 435         test()
 436     except KeyboardInterrupt:
 437         print "\n[interrupt]"