From 01ecf2bfb1ffb72159a46f049170ac2ca98deb2b Mon Sep 17 00:00:00 2001
From: Lars Wirzenius <liw@liw.fi>
Date: Mon, 18 Apr 2016 20:07:31 +0300
Subject: Add SRCDIR to PYTHONPATH if shell is Python

---
 yarn | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

(limited to 'yarn')

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))
 
-- 
cgit v1.2.1