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