Added @file feature. Masked FFTW 3D tests.

3 files changed, 53 insertions, 7 deletions

diff --git a/numbench/benchutils.py b/numbench/benchutils.py
index df12ee9..7f2542d 100644
--- a/numbench/benchutils.py
+++ b/numbench/benchutils.py
@@ -1,8 +1,18 @@
-import os, sys
+import os, sys, string, random
 import subprocess as sp
 
+__all__ = ['mkdir', 'tmpfile', 'run_cmd']
+
 def mkdir(dir):
     if not os.path.exists(dir):
         os.makedirs(dir)
-        
+    
+def tmpfile(dir="/var/tmp"):
+    """Returns the path of a free temporary file within the given directory."""
+    chars = string.letters + string.digits
+    while True:
+        fname = os.path.join(dir, random.sample(chars, 10))
+        if not os.path.exists(fname):
+            return fname
+    
 run_cmd = lambda c : sp.Popen(c, stdout=sp.PIPE).communicate()[0]
diff --git a/numbench/fftw.py b/numbench/fftw.py
index e462716..1596921 100644
--- a/numbench/fftw.py
+++ b/numbench/fftw.py
@@ -10,11 +10,12 @@ class Module(btlbase.BTLBase):
           "FFTW_2D_Forward_Measure", "FFTW_2D_Forward_Estimate",
           "FFTW_2D_Backward_Measure", "FFTW_2D_Backward_Estimate",
           
-          "FFTW_3D_Forward_Measure", "FFTW_3D_Forward_Estimate",
-          "FFTW_3D_Backward_Measure", "FFTW_3D_Backward_Estimate"
+          # Mask 3D tests while unstable
+#          "FFTW_3D_Forward_Measure", "FFTW_3D_Forward_Estimate",
+#          "FFTW_3D_Backward_Measure", "FFTW_3D_Backward_Estimate"
         )
     
-    def _parse_args(self, args):     
+    def _parse_args(self, args):
         # Parse arguments
         tests = []
         for i in args:
diff --git a/numbench/main.py b/numbench/main.py
index 6150611..58e5e58 100755
--- a/numbench/main.py
+++ b/numbench/main.py
@@ -1,6 +1,6 @@
 #! /usr/bin/env python2
 
-import os, sys, signal, shlex, time
+import os, sys, signal, shlex, shutil, time
 from os.path import join as pjoin
 import subprocess as sp
 
@@ -39,6 +39,20 @@ def print_help():
     print "More information about a module is available through the command:"
     print "  numbench module --help"
     
+def readEnvFile(fname):
+    """Reads a bash file with void environment and returns the environment
+    at the end of the execution."""
+    proc = sp.Popen('. '+fname+' &> /dev/null; env', \
+      shell=True, stdout=sp.PIPE, env={})
+    lines = proc.stdout.read().split('\n')[:-1]
+    env = dict([l.split('=', 1) for l in lines])
+    
+    for k in ('SHLVL', 'PWD', '_'):
+        if env.has_key(k):
+            del env[k]
+    return env
+    
+    
 def tests_from_input(input):
     tests = {}
     for line in input.split('\n'):
@@ -49,26 +63,46 @@ def tests_from_input(input):
         if line[0] == '#':
             continue
         env = {}
-        # TODO: add @file for env set based on external file
         skip = []
         change = {}
         descr = None
+        fileenv = {}
+        
+        # Interpret arguments
         for var in spl[2:]:
+            
+            # if begins with '-': skip implementation
             if var[0] == '-':
                 skip.append(var[1:])
+                
+            # if key:value, substitute pkg-config dependency
             elif ':' in var and not '=' in var:
                 c_0, c_1 = var.split(':', 1)
                 change[c_0] = c_1
+                
+            # if descr|text set description (for future use)
             elif var[:6] == 'descr|':
                 descr = var[6:]
+                
+            # if @file: read bash script and set env
+            elif var[0] == '@':
+                fileenv = readEnvFile(pjoin(cfg.curdir, var[1:]))
+            
+            # Otherwise, assume key=value syntax
             else:
                 e_0, e_1 = var.split('=', 1)
                 env[e_0] = e_1
+                
+        # Set environment (argument overrides bash file)
+        env = dict( fileenv.items() + env.items() )
+        
         try:
+            # Insert test
             avail = available_packages(spl[1])[-1]
             tests[spl[0]] = {'package':avail , 'env':env, 'skip':skip, \
               'changes':change, 'descr':descr}
         except:
+            # Or trigger an non-fatal error
             sys.stderr.write('Error: package ' + spl[1] + ' not found\n')
     return tests
 
@@ -79,6 +113,7 @@ def tests_from_input(input):
 
 import benchconfig as cfg
 import benchchilds
+import benchutils as bu
 
 
 # If no argument is given, print the help