pyabc changes for HWMCC13

5 files changed, 150 insertions, 46 deletions

diff --git a/src/python/module.make b/src/python/module.make
index 09e2de3f..51be60af 100644
--- a/src/python/module.make
+++ b/src/python/module.make
@@ -34,6 +34,12 @@ ifdef ABC_PYTHON
         $(ABC_PYTHON_FILES_PREFIX)/dist \
         pyabc.tgz
 
+    ABC_PYABC_DIR ?= pyabc
+    ABC_PYABC_TGZ ?= pyabc.tgz
+    ABC_PYABC_EXTRA_BIN ?=
+    ABC_PYABC_EXTRA_LIB ?=
+
+
 %_wrap.c %.py : %.i
     $(ABC_SWIG) -python -outdir $(<D) $<
 
@@ -53,11 +59,18 @@ pyabc_extension_install : pyabc_extension_build
 pyabc_extension_bdist : pyabc_extension_build    
     ( cd $(ABC_PYTHON_FILES_PREFIX) && python setup.py bdist )
     
-pyabc.tgz : $(PROG) $(ABC_PYTHON_SRC:_wrap.c=.py) $(ABC_PYTHON_FILES_PREFIX)/abc.sh $(ABC_PYTHON_FILES_PREFIX)/package.py
+.PHONY: pyabc_tgz
+
+pyabc_tgz : $(ABC_PYABC_TGZ)
+    
+$(ABC_PYABC_TGZ) : $(PROG) $(ABC_PYTHON_SRC:_wrap.c=.py) $(ABC_PYTHON_FILES_PREFIX)/abc.sh $(ABC_PYTHON_FILES_PREFIX)/package.py
     $(ABC_PYTHON) $(ABC_PYTHON_FILES_PREFIX)/package.py \
+            --pyabc_dir=$(ABC_PYABC_DIR) \
         --abc=$(PROG) \
         --abc_sh=$(ABC_PYTHON_FILES_PREFIX)/abc.sh \
         --pyabc=$(ABC_PYTHON_FILES_PREFIX) \
+        --extra_bin="$(ABC_PYABC_EXTRA_BIN)" \
+        --extra_lib="$(ABC_PYABC_EXTRA_LIB)" \
         --out=$@ \
         $(ABC_PYTHON_OPTIONS)
 
diff --git a/src/python/package.py b/src/python/package.py
index 1764ab6f..ab464e0a 100644
--- a/src/python/package.py
+++ b/src/python/package.py
@@ -5,6 +5,7 @@ import zipfile
 import tarfile
 import tempfile
 import time
+import py_compile
 
 def zip_library(f, extra_files = []):
     lib = "%s/lib/python%s/"%(sys.prefix,sys.version[:3])
@@ -15,13 +16,25 @@ def zip_library(f, extra_files = []):
         arcroot = os.path.relpath(root, lib)
         for f in files:
             _, ext = os.path.splitext(f)
-            if ext in ['.py', '.pyo', '.pyc']:
+            if ext in ['.py']:
                 zf.write(os.path.join(root,f), os.path.join(arcroot, f))
 
     for s, r in extra_files:
         zf.write( s, r )
 
     zf.close()
+    
+def add_python_lib(tf, lib_dir, lib, mtime):
+    
+    for root, _, files in os.walk(lib):
+        
+        arcroot = os.path.join( lib_dir, os.path.relpath(root, lib) )
+        add_dir(tf, arcroot, mtime)
+        
+        for f in files:
+            _, ext = os.path.splitext(f)
+            if ext in ['.py']:
+                add_file( tf, os.path.join(root,f), os.path.join(arcroot, f), 0666, mtime)
 
 def add_dir(tf, dir, mtime):
     ti = tarfile.TarInfo(dir)
@@ -43,21 +56,23 @@ def add_fileobj(tf, f, arcname, mode, mtime):
     tf.addfile(ti, f)
     
 def add_file(tf, fname, arcname, mode, mtime):
-    f = open(fname, "rb")
-    add_fileobj(tf, f, arcname, mode, mtime)
-    f.close()
+    print "\t adding %s as %s"%(fname, arcname)
+    
+    with open(fname, "rb") as f:
+        add_fileobj(tf, f, arcname, mode, mtime)
 
-def package(abc_exe, abc_sh, pyabc, ofname, scripts_dir, use_sys):
+def package(pyabc_dir, extra_bin, extra_lib, abc_exe, abc_sh, pyabc, ofname, scripts_dir, use_sys):
+    
     mtime = time.time()
     
     tf = tarfile.open(ofname, "w:gz")
     
-    add_dir(tf, "pyabc", mtime)
+    add_dir(tf, "%s"%pyabc_dir, mtime)
     
-    add_dir(tf, "pyabc/bin", mtime)
+    add_dir(tf, "%s/bin"%pyabc_dir, mtime)
     
-    add_file(tf, abc_exe, "pyabc/bin/abc_exe", 0777, mtime)
-    add_file(tf, abc_sh, "pyabc/bin/abc", 0777, mtime)
+    add_file(tf, abc_exe, "%s/bin/abc_exe"%pyabc_dir, 0777, mtime)
+    add_file(tf, abc_sh, "%s/bin/abc"%pyabc_dir, 0777, mtime)
 
     if scripts_dir:
         for fn in os.listdir(scripts_dir):
@@ -65,15 +80,23 @@ def package(abc_exe, abc_sh, pyabc, ofname, scripts_dir, use_sys):
             if os.path.isfile(fullname):
                 fnroot, fnext = os.path.splitext(fn)
                 if fnext==".sh":
-                    add_file( tf, fullname, os.path.join("pyabc/bin", fnroot), 0777, mtime)
-                else:
-                    add_file( tf, fullname, os.path.join("pyabc/scripts", fn), 0666, mtime)
+                    add_file( tf, fullname, os.path.join("%s/bin"%pyabc_dir, fnroot), 0777, mtime)
+                elif fnext not in ( '.pyc', '.pyo'):
+                    add_file( tf, fullname, os.path.join("%s/scripts"%pyabc_dir, fn), 0666, mtime)
     
-    add_dir(tf, "pyabc/lib", mtime)
+    for bin in extra_bin:
+        add_file( tf, bin, os.path.join("%s/bin"%pyabc_dir, os.path.basename(bin)), 0777, mtime)
+        
+    lib_dir = "%s/lib"%pyabc_dir
+
+    add_dir(tf, lib_dir, mtime)
+
+    for lib in extra_lib:
+        add_python_lib( tf, lib_dir, lib, mtime)
     
     for entry in os.listdir(pyabc):
         if entry.endswith('.py'):
-            add_file( tf, os.path.join(pyabc, entry), os.path.join("pyabc/lib", entry), 0666, mtime)
+            add_file( tf, os.path.join(pyabc, entry), os.path.join("%s/lib"%pyabc_dir, entry), 0666, mtime)
     
     if not use_sys:
         # ZIP standard library    
@@ -82,7 +105,7 @@ def package(abc_exe, abc_sh, pyabc, ofname, scripts_dir, use_sys):
         zip_library(zf, [])
         zf.flush()
         
-        add_fileobj(tf, zf, "pyabc/lib/python_library.zip", 0666, mtime)
+        add_fileobj(tf, zf, "%s/lib/python_library.zip"%pyabc_dir, 0666, mtime)
         
         zf.close()
     
@@ -93,7 +116,7 @@ def package(abc_exe, abc_sh, pyabc, ofname, scripts_dir, use_sys):
         for fn in os.listdir(lib_dynload):
             fullname = os.path.join(lib_dynload, fn)
             if os.path.isfile(fullname):
-                add_file( tf, fullname, os.path.join("pyabc/lib", fn), 0666, mtime)
+                add_file( tf, fullname, os.path.join("%s/lib"%pyabc_dir, fn), 0666, mtime)
     
     tf.close()
 
@@ -104,6 +127,9 @@ def main(args):
 
     parser = optparse.OptionParser(usage)
 
+    parser.add_option("-d", "--pyabc_dir", dest="pyabc_dir", help="name of generated directory" )
+    parser.add_option("-b", "--extra_bin", dest="extra_bin", help="extra binaries to pack" )
+    parser.add_option("-l", "--extra_lib", dest="extra_lib", help="extra directories in lib to pack" )
     parser.add_option("-a", "--abc", dest="abc", help="location of the ABC exeutable")
     parser.add_option("-s", "--abc_sh", dest="abc_sh", help="location of the ABC setup script")
     parser.add_option("-p", "--pyabc", dest="pyabc", help="location of pyabc.py")
@@ -117,11 +143,14 @@ def main(args):
         parser.print_help()
         return 1
         
-    if not options.abc or not options.abc_sh or not options.pyabc or not options.out:
+    if not options.pyabc_dir or not options.abc or not options.abc_sh or not options.pyabc or not options.out:
         parser.print_help()
         return 1
 
-    return package(options.abc, options.abc_sh, options.pyabc, options.out, options.scripts, options.sys)
+    extra_bin = options.extra_bin.split(',') if options.extra_bin else []
+    extra_lib = options.extra_lib.split(',') if options.extra_lib else []
+
+    return package(options.pyabc_dir, extra_bin, extra_lib, options.abc, options.abc_sh, options.pyabc, options.out, options.scripts, options.sys)
 
 if __name__=="__main__":
     main(sys.argv)
diff --git a/src/python/pyabc.i b/src/python/pyabc.i
index 1d301c89..51b36f20 100644
--- a/src/python/pyabc.i
+++ b/src/python/pyabc.i
@@ -430,7 +430,7 @@ void pyabc_internal_register_command( char * sGroup, char * sName, int fChanges
 {
     Abc_Frame_t* pAbc = Abc_FrameGetGlobalFrame();
 
-    Cmd_CommandAdd( pAbc, sGroup, sName, (int(*)(Abc_Frame_t*, int, char**))pyabc_internal_abc_command_callback, fChanges);
+    Cmd_CommandAdd( pAbc, sGroup, sName, (Cmd_CommandFuncType)pyabc_internal_abc_command_callback, fChanges);
 }
 
 static int sigchld_pipe_fd = -1;
@@ -637,7 +637,14 @@ int _posix_kill(int pid, int signum)
 
 void _set_death_signal()
 {
+    // send SIGINT if parent process is dead
     prctl(PR_SET_PDEATHSIG, SIGINT);
+    
+    // if parent process is already dead (and adopted by init)
+    if ( getppid() == 1)
+    {
+        raise(SIGINT);
+    }
 }
 
 %}
diff --git a/src/python/pyabc_split.py b/src/python/pyabc_split.py
index bbdcb249..89038a2c 100644
--- a/src/python/pyabc_split.py
+++ b/src/python/pyabc_split.py
@@ -90,6 +90,8 @@ import cPickle as pickle
 import signal
 import cStringIO
 
+import traceback
+
 from contextlib import contextmanager
 
 import pyabc
@@ -107,8 +109,7 @@ def _retry_select(rlist):
 
 class _splitter(object):
     
-    def __init__(self, funcs):
-        self.funcs = funcs
+    def __init__(self):
         self.pids = []
         self.fds = {}
         self.buffers = {}
@@ -117,27 +118,49 @@ class _splitter(object):
     def is_done(self):
         return len(self.fds) == 0
     
-    def cleanup(self):
+    def _kill(self, pids):
+
         # close pipes and kill child processes
-        for pid,(i,fd) in self.fds.iteritems():
+        for pid in pids:
+
+            if pid == -1:
+                continue
+
+            i, fd = self.fds[pid]
+            
+            del self.buffers[fd]
+            del self.fds[pid]
+
+            self.pids[i] = -1
+            self.results[pid] = None
+            
             os.close(fd)
+            
             try:
-                os.kill( pid, signal.SIGINT )
+                os.kill( pid, signal.SIGINT)
             except Exception as e:
                 print >>sys.stderr, 'exception while trying to kill pid=%d: '%pid, e
                 raise
-        
-            # wait for termination and update result
-        for pid, _ in self.fds.iteritems():
+
+        # wait for termination and update result
+        for pid in pids:
             os.waitpid( pid, 0 )
-            self.results[pid] = None
+
+    def kill(self, ids):
+        
+        self._kill( [ self.pids[i] for i in ids ] )
             
-        self.fds = {}
-        self.buffers = {}
+    def cleanup(self):
+        self._kill( self.fds.keys() )
 
     def child( self, fdw, f):
+
         # call function
-        res = f()
+        try:
+            res = f()
+        except:
+            traceback.print_exc()
+            raise
         
         # write return value into pipe
         with os.fdopen( fdw, "w" ) as fout:
@@ -145,7 +168,7 @@ class _splitter(object):
             
         return 0
 
-    def fork_one(self, f):
+    def _fork_one(self, f):
         
         # create a pipe to communicate with the child process
         pr,pw = os.pipe()
@@ -177,13 +200,17 @@ class _splitter(object):
             if os.getpid() != parentpid:
                 os._exit(rc)
 
-    def fork_all(self):
-        for i,f in enumerate(self.funcs):
-            pid, fd = self.fork_one(f)
-            self.pids.append(pid)
-            self.fds[pid] = (i,fd)
-            self.buffers[fd] = cStringIO.StringIO()
+    def fork_one(self, func):
+        pid, fd = self._fork_one(func)
+        i = len(self.pids)
+        self.pids.append(pid)
+        self.fds[pid] = (i, fd)
+        self.buffers[fd] = cStringIO.StringIO()
+        return i
     
+    def fork_all(self, funcs):
+        return [ self.fork_one(f) for f in funcs ]
+
     def communicate(self):
         
         rlist = [ fd for _, (_,fd) in self.fds.iteritems() ]
@@ -216,6 +243,9 @@ class _splitter(object):
         i, fd = self.fds[pid]
         del self.fds[pid]
 
+        # remove the pid
+        self.pids[i] = -1
+
         # retrieve the buffer
         buffer = self.buffers[fd]
         del self.buffers[fd]
@@ -226,29 +256,37 @@ class _splitter(object):
             if not s:
                 break
             buffer.write(s)
+            
+        os.close(fd)
 
         try:
             return (i, pickle.loads(buffer.getvalue()))
         except EOFError, pickle.UnpicklingError:
             return (i, None)
+            
+    def __iter__(self):
+        def iterator():
+            while not self.is_done():
+                yield self.get_next_result()
+        return iterator()
 
 @contextmanager
-def _splitter_wrapper(funcs):
+def make_splitter():
     # ensure cleanup of child processes
-    s = _splitter(funcs)
+    s = _splitter()
     try:
         yield s
     finally:
         s.cleanup()
-
+        
 def split_all_full(funcs):
     # provide an iterator for child process result
-    with _splitter_wrapper(funcs) as s:
+    with make_splitter() as s:
         
-        s.fork_all()
+        s.fork_all(funcs)
         
-        while not s.is_done():
-            yield s.get_next_result()
+        for res in s:
+            yield res
 
 def defer(f):
     return lambda *args, **kwargs: lambda : f(*args,**kwargs)
diff --git a/src/python/redirect.py b/src/python/redirect.py
index 498fe150..0afccb77 100644
--- a/src/python/redirect.py
+++ b/src/python/redirect.py
@@ -25,6 +25,23 @@ def _dup( f ):
     os.close(fd)
     
 @contextmanager
+def save_stdout( src = sys.stdout ):
+    """
+    Redirect 
+    """
+    fd = os.dup( src.fileno() )
+    own = True
+    
+    try:
+        with os.fdopen( fd, "w", 0) as f:
+            own = False
+            yield f
+    except:
+        if own:
+            os.close(fd)
+        raise
+    
+@contextmanager
 def redirect(dst = null_file, src = sys.stdout):
     
     """