Add SRCDIR to PYTHONPATH if shell is Python

1 files changed, 15 insertions, 0 deletions

diff --git a/yarn b/yarn
index 7e492b4..935640c 100755
--- a/yarn
+++ b/yarn
@@ -463,6 +463,7 @@ class YarnRunner(cliapp.Application):
         env['HOME'] = self.homedir(datadir)
         for i, match in enumerate(m.groups('')):
             env['MATCH_%d' % (i+1)] = match.encode('utf-8')
+        self.add_srcdir_to_pythonpath(env, env['SRCDIR'])
 
         if self.settings['cd-datadir']:
             cwd = datadir
@@ -509,6 +510,20 @@ class YarnRunner(cliapp.Application):
 
         return exit
 
+    def add_srcdir_to_pythonpath(self, env, srcdir):
+        # Special handling of PYTHONPATH. Add $SRCDIR to it so that
+        # Python IMPLEMENTS can use it. But only if --shell=python is
+        # used. This is a bit of magic that hopefully won't surprise
+        # users too much.
+
+        if self.settings['shell'] == 'python':
+            pythonpath = env.get('PYTHONPATH', None)
+            if pythonpath:
+                pythonpath += ':' + srcdir
+            else:
+                pythonpath = srcdir
+            env['PYTHONPATH'] = pythonpath
+
     def scenario_dir(self, tempdir, scenario):
         return os.path.join(tempdir, self.nice(scenario.name))