import os
from pyspark import SparkContext, SparkConf, SparkFiles

class SparkContextVenv(SparkContext):
    def __init__(self,app_name, virtual_env_tarball_file, venv_name=None, python_path='/usr/bin/python2.7', env_variables_dict=None):
        assert os.path.isfile(virtual_env_tarball_file)
        assert os.path.isfile(python_path)
        self.virtual_env_tarball_file = virtual_env_tarball_file
        self.python_path = python_path
        self.sc = None
        if not venv_name:
            self.venv_name = os.path.basename(virtual_env_tarball_file).split('.')[0] #this assumes that the name of the tarball is the same of the folder stored inside
        else:
            self.venv_name = venv_name

        os.environ['PYSPARK_PYTHON'] = self.python_path
        conf = SparkConf()
        self.env_variables_dict = env_variables_dict
        if env_variables_dict:
            for key, value in env_variables_dict.iteritems():
                print 'setting env: {} with value: {}'.format(key, value)
                conf.setExecutorEnv(key=key, value=value)
        super(SparkContextVenv, self).__init__(appName=app_name, conf=conf)
        self.addFile(self.virtual_env_tarball_file)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.stop()

def virtualenv(venv_name):
  def wrap(f):
    def wrapped_f(*args, **kwargs):
        venv_location = SparkFiles.get(venv_name)
        print 'called wrapped_f for virtual env in {}'.format(venv_location)
        activate_env="%s/bin/activate_this.py" % venv_location
        execfile(activate_env, dict(__file__=activate_env))
        return f(*args, **kwargs)
    return wrapped_f
  return wrap


@@ -154,16 +154,21 @@
 all_file
+=None
 , venv_n
@@ -237,16 +237,17 @@
 =None):%0A
+%0A
         
@@ -268,16 +268,41 @@
 .isfile(
+python_path)%0A%0A        if 
 virtual_
@@ -317,26 +317,27 @@
 all_file
-)%0A
+:%0A 
         
 assert o
@@ -320,32 +320,35 @@
 _file:%0A         
+   
 assert os.path.i
@@ -353,29 +353,46 @@
 .isfile(
-python_path)%0A
+virtual_env_tarball_file)%0A    
         
@@ -456,66 +456,8 @@
     
-    self.python_path = python_path%0A        self.sc = None%0A
     
@@ -478,16 +478,20 @@
 v_name:%0A
+    
         
@@ -663,14 +663,22 @@
     
+    
 else:%0A
+    
     
@@ -711,16 +711,177 @@
 env_name
+%0A        else:%0A            self.virtual_env_tarball_file = None%0A            self.venv_name = None%0A%0A%0A        self.python_path = python_path%0A        self.sc = None
 %0A%0A      
@@ -1311,16 +1311,62 @@
 f=conf)%0A
+        if self.virtual_env_tarball_file:%0A    
         
@@ -1856,16 +1856,193 @@
 e_env))%0A
+%0A        import sys%0A        numpy_mods = %5Bitem for item in sys.modules if ('numpy' in item or 'scipy' in item)%5D%0A        for mod in numpy_mods:%0A            sys.modules.pop(mod)%0A%0A
         

fix numpy problem