daily pastebin goal
54%
SHARE
TWEET

Untitled

a guest Feb 17th, 2019 69 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. # coding: utf-8
  2.  
  3. from subprocess import Popen
  4. import subprocess
  5. import win32job
  6. import win32process
  7. import win32api
  8.  
  9.  
  10. class JobPopen(Popen):
  11.     """Start a process in a new Win32 job object.
  12.    
  13.     This `subprocess.Popen` subclass takes the same arguments as Popen and
  14.     behaves the same way. In addition to that, created processes will be
  15.     assigned to a new anonymous Win32 job object on startup, which will
  16.     guarantee that the processes will be terminated by the OS as soon as
  17.     either the Popen object, job object handle or parent Python process are
  18.     closed.
  19.     """
  20.  
  21.     class _winapijobhandler(object):
  22.         """Patches the native CreateProcess function in the subprocess module
  23.         to assign created threads to the given job"""
  24.  
  25.         def __init__(self, oldapi, job):
  26.             self._oldapi = oldapi
  27.             self._job = job
  28.  
  29.         def __getattr__(self, key):
  30.             if key != "CreateProcess":
  31.                 return getattr(self._oldapi, key)  # Any other function is run as before
  32.             else:
  33.                 return self.CreateProcess  # CreateProcess will call the function below
  34.  
  35.         def CreateProcess(self, *args, **kwargs):
  36.             hp, ht, pid, tid = self._oldapi.CreateProcess(*args, **kwargs)
  37.             win32job.AssignProcessToJobObject(self._job, hp)
  38.             win32process.ResumeThread(ht)
  39.             return hp, ht, pid, tid
  40.  
  41.     def __init__(self, *args, **kwargs):
  42.         """Start a new process using an anonymous job object. Takes the same arguments as Popen"""
  43.  
  44.         # Create a new job object
  45.         self._win32_job = self._create_job_object()
  46.  
  47.         # Temporarily patch the subprocess creation logic to assign created
  48.         # processes to the new job, then resume execution normally.
  49.         CREATE_SUSPENDED = 0x00000004
  50.         kwargs.setdefault("creationflags", 0)
  51.         kwargs["creationflags"] |= CREATE_SUSPENDED
  52.         try:
  53.             _winapi = subprocess._winapi  # Python 3
  54.             _winapi_key = "_winapi"
  55.         except AttributeError:
  56.             _winapi = subprocess._subprocess  # Python 2
  57.             _winapi_key = "_subprocess"
  58.         try:
  59.             setattr(subprocess, _winapi_key, JobPopen._winapijobhandler(_winapi, self._win32_job))
  60.             super(JobPopen, self).__init__(*args, **kwargs)
  61.         finally:
  62.             setattr(subprocess, _winapi_key, _winapi)
  63.  
  64.     def _create_job_object(self):
  65.         """Create a new anonymous job object"""
  66.         hjob = win32job.CreateJobObject(None, "")
  67.         extended_info = win32job.QueryInformationJobObject(hjob, win32job.JobObjectExtendedLimitInformation)
  68.         extended_info['BasicLimitInformation']['LimitFlags'] = win32job.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
  69.         win32job.SetInformationJobObject(hjob, win32job.JobObjectExtendedLimitInformation, extended_info)
  70.         return hjob
  71.  
  72.     def _close_job_object(self, hjob):
  73.         """Close the handle to a job object, terminating all processes inside it"""
  74.         if self._win32_job:
  75.             win32api.CloseHandle(self._win32_job)
  76.             self._win32_job = None
  77.  
  78.     # This ensures that no remaining subprocesses are found when the process
  79.     # exits from a `with JobPopen(...)` block.
  80.     def __exit__(self, exc_type, value, traceback):
  81.         super(JobPopen, self).__exit__(exc_type, value, traceback)
  82.         self._close_job_object(self._win32_job)
  83.  
  84.     # Python does not keep a reference outside of the parent class when the
  85.     # interpreter exits, which is why we keep it here.
  86.     _Popen = subprocess.Popen  
  87.     def __del__(self):
  88.         self._Popen.__del__(self)
  89.         self._close_job_object(self._win32_job)
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
 
Top