Add SRCDIR to PYTHONPATH if shell is Python

2 files changed, 17 insertions, 0 deletions

diff --git a/NEWS b/NEWS
index ca4c92a..878ebf9 100644
--- a/NEWS
+++ b/NEWS
@@ -9,6 +9,8 @@ Version 0.22+git, not yet released
 * Yarn now has the `--stop-on-first-fail` setting, which makes it stop
   if a scenario step fails, and not continue with the next scenario.
 
+* Yarn now adds SRCDIR to PYTHONPATH, if `--shell=python` is used.
+
 Version 0.22, released 2016-02-04
 ---------------------------------
 
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))