hooke/compat/forking.py

   1 # Copyright (C) 2010 W. Trevor King <wking@drexel.edu>
   2 #
   3 # This file is part of Hooke.
   4 #
   5 # Hooke is free software: you can redistribute it and/or modify it
   6 # under the terms of the GNU Lesser General Public License as
   7 # published by the Free Software Foundation, either version 3 of the
   8 # License, or (at your option) any later version.
   9 #
  10 # Hooke is distributed in the hope that it will be useful, but WITHOUT
  11 # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
  12 # or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General
  13 # Public License for more details.
  14 #
  15 # You should have received a copy of the GNU Lesser General Public
  16 # License along with Hooke.  If not, see
  17 # <http://www.gnu.org/licenses/>.
  18
  19 """Dynamically patch :func:`multiprocessing.forking.prepare`.
  20
  21 `prepare` can have trouble `loading modules from .eggs`_ or
  22 `finding hooke on Windows`_.  Importing this module applies the
  23 suggested patch dynamically.
  24
  25 .. _loading modules from .eggs: http://bugs.python.org/issue8534
  26 .. _finding hooke on Windows:
  27   http://code.google.com/p/hooke/issues/detail?id=40#c17
  28 """
  29
  30 import multiprocessing.forking
  31
  32
  33 def prepare(data):
  34     '''
  35     Try to get current process ready to unpickle process object
  36     '''
  37     old_main_modules.append(sys.modules['__main__'])
  38
  39     if 'name' in data:
  40         process.current_process().name = data['name']
  41
  42     if 'authkey' in data:
  43         process.current_process()._authkey = data['authkey']
  44
  45     if 'log_to_stderr' in data and data['log_to_stderr']:
  46         util.log_to_stderr()
  47
  48     if 'log_level' in data:
  49         util.get_logger().setLevel(data['log_level'])
  50
  51     if 'sys_path' in data:
  52         sys.path = data['sys_path']
  53
  54     if 'sys_argv' in data:
  55         sys.argv = data['sys_argv']
  56
  57     if 'dir' in data:
  58         os.chdir(data['dir'])
  59
  60     if 'orig_dir' in data:
  61         process.ORIGINAL_DIR = data['orig_dir']
  62
  63     if 'main_path' in data:
  64         main_path = data['main_path']
  65         main_name = os.path.splitext(os.path.basename(main_path))[0]
  66         if main_name == '__init__':
  67             main_name = os.path.basename(os.path.dirname(main_path))
  68
  69         if main_name != 'ipython':
  70             import imp
  71
  72             if main_path is None:
  73                 dirs = None
  74             elif os.path.basename(main_path).startswith('__init__.py'):
  75                 dirs = [os.path.dirname(os.path.dirname(main_path))]
  76             else:
  77                 dirs = [os.path.dirname(main_path)]
  78
  79             assert main_name not in sys.modules, main_name
  80
  81             try:
  82                 main_module = __import__(main_name)
  83             except ImportError:
  84                 file, path_name, etc = imp.find_module(main_name, dirs)
  85
  86                 try:
  87                     # We would like to do "imp.load_module('__main__', ...)"
  88                     # here.  However, that would cause 'if __name__ ==
  89                     # "__main__"' clauses to be executed.
  90                     main_module = imp.load_module(
  91                         '__parents_main__', file, path_name, etc
  92                         )
  93                 finally:
  94                     if file:
  95                         file.close()
  96
  97             sys.modules['__main__'] = main_module
  98             main_module.__name__ = '__main__'
  99
 100             # Try to make the potentially picklable objects in
 101             # sys.modules['__main__'] realize they are in the main
 102             # module -- somewhat ugly.
 103             for obj in main_module.__dict__.values():
 104                 try:
 105                     if obj.__module__ == '__parents_main__':
 106                         obj.__module__ = '__main__'
 107                 except Exception:
 108                     pass
 109 multiprocessing.forking.prepare = prepare