write the __main__ module to __disco__.py to be safe and make writemo…

…dule more like write

lib/disco/fileutils.py

@@ -1,6 +1,7 @@
 import os, struct, sys, time
 from cPickle import dumps
 from cStringIO import StringIO
+from inspect import getmodule, getsourcefile
 from zipfile import ZipFile, ZIP_DEFLATED
 from zlib import compress, crc32
 
@@ -139,14 +140,12 @@ def writepath(self, pathname, exclude=()):
             if ext not in exclude:
                 self.write(file, file)
 
-    def writemodule(self, path, module):
-        from inspect import getsource
+    def writemodule(self, module, arcname=None):
         if isinstance(module, basestring):
             module = __import__(module)
-        self.writestr(path, getsource(module))
+        self.write(getsourcefile(module), arcname=arcname)
 
     def writesource(self, object):
-        from inspect import getmodule, getsourcefile
         self.writepath(getsourcefile(getmodule(object)))
 
     def dump(self, handle):

lib/disco/worker/__init__.py

@@ -239,7 +239,7 @@ def jobzip(self, job, **jobargs):
         jobzip.writepath(os.path.dirname(discopath), exclude=('.pyc',))
         jobzip.writesource(job)
         jobzip.writesource(self)
-        jobzip.writemodule('lib/__main__.py', '__main__')
+        jobzip.writemodule('__main__', '__disco__.py')
         return jobzip
 
     def input(self, task, merged=False, **kwds):
@@ -327,7 +327,7 @@ def main(cls):
             cls.send('PID', os.getpid())
             try:
                 from imp import find_module, load_module
-                __disco__ = load_module('__disco__', *find_module('__main__', ['lib']))
+                __disco__ = load_module('__disco__', *find_module('__disco__', ['']))
                 sys.modules['__main__'].__dict__.update(__disco__.__dict__)
             except ImportError:
                 pass