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 logging
  31 import multiprocessing.forking
  32
  33
  34 def prepare(data):
  35     '''
  36     Try to get current process ready to unpickle process object
  37     '''
  38     old_main_modules.append(sys.modules['__main__'])
  39
  40     if 'name' in data:
  41         process.current_process().name = data['name']
  42
  43     if 'authkey' in data:
  44         process.current_process()._authkey = data['authkey']
  45
  46     if 'log_to_stderr' in data and data['log_to_stderr']:
  47         util.log_to_stderr()
  48
  49     if 'log_level' in data:
  50         util.get_logger().setLevel(data['log_level'])
  51
  52     if 'sys_path' in data:
  53         sys.path = data['sys_path']
  54
  55     if 'sys_argv' in data:
  56         sys.argv = data['sys_argv']
  57
  58     if 'dir' in data:
  59         os.chdir(data['dir'])
  60
  61     if 'orig_dir' in data:
  62         process.ORIGINAL_DIR = data['orig_dir']
  63
  64     if 'main_path' in data:
  65         main_path = data['main_path']
  66         main_name = os.path.splitext(os.path.basename(main_path))[0]
  67         if main_name == '__init__':
  68             main_name = os.path.basename(os.path.dirname(main_path))
  69
  70         if main_name != 'ipython':
  71             import imp
  72
  73             if main_path is None:
  74                 dirs = None
  75             elif os.path.basename(main_path).startswith('__init__.py'):
  76                 dirs = [os.path.dirname(os.path.dirname(main_path))]
  77             else:
  78                 dirs = [os.path.dirname(main_path)]
  79
  80             assert main_name not in sys.modules, main_name
  81
  82             try:
  83                 main_module = __import__(main_name)
  84             except ImportError:
  85                 file, path_name, etc = imp.find_module(main_name, dirs)
  86
  87                 try:
  88                     # We would like to do "imp.load_module('__main__', ...)"
  89                     # here.  However, that would cause 'if __name__ ==
  90                     # "__main__"' clauses to be executed.
  91                     main_module = imp.load_module(
  92                         '__parents_main__', file, path_name, etc
  93                         )
  94                 finally:
  95                     if file:
  96                         file.close()
  97
  98             sys.modules['__main__'] = main_module
  99             main_module.__name__ = '__main__'
 100
 101             # Try to make the potentially picklable objects in
 102             # sys.modules['__main__'] realize they are in the main
 103             # module -- somewhat ugly.
 104             for obj in main_module.__dict__.values():
 105                 try:
 106                     if obj.__module__ == '__parents_main__':
 107                         obj.__module__ = '__main__'
 108                 except Exception:
 109                     pass
 110 multiprocessing.forking.prepare = prepare
 111
 112 logging.warn(
 113     'Monkey patched multiprocessing.forking.prepare for issue8534')