From e8af3c8e9870be6304fb835b598184c8651f908a Mon Sep 17 00:00:00 2001
From: Maista <131594657+Maista6969@users.noreply.github.com>
Date: Wed, 20 Dec 2023 03:32:19 +0100
Subject: [PATCH] Set PYTHONPATH environment variable for Python script
 scrapers (#4372)

* Set PYTHONPATH environment variable for Python script scrapers
* Convert PYTHONPATH to absolute
* Generalise and apply to plugins

---------

Co-authored-by: WithoutPants <53250216+WithoutPants@users.noreply.github.com>
---
 pkg/plugin/raw.go     |  4 ++++
 pkg/python/env.go     | 15 +++++++++++++++
 pkg/scraper/script.go |  2 ++
 3 files changed, 21 insertions(+)
 create mode 100644 pkg/python/env.go

diff --git a/pkg/plugin/raw.go b/pkg/plugin/raw.go
index 5c2fcea4e..f276e139c 100644
--- a/pkg/plugin/raw.go
+++ b/pkg/plugin/raw.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"io"
 	"os/exec"
+	"path/filepath"
 	"strings"
 	"sync"
 
@@ -52,6 +53,9 @@ func (t *rawPluginTask) Start() error {
 			logger.Warnf("%s", err)
 		} else {
 			cmd = p.Command(context.TODO(), command[1:])
+
+			envVariable, _ := filepath.Abs(filepath.Dir(filepath.Dir(t.plugin.path)))
+			python.AppendPythonPath(cmd, envVariable)
 		}
 	}
 
diff --git a/pkg/python/env.go b/pkg/python/env.go
new file mode 100644
index 000000000..004c1ad91
--- /dev/null
+++ b/pkg/python/env.go
@@ -0,0 +1,15 @@
+package python
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+)
+
+func AppendPythonPath(cmd *exec.Cmd, path string) {
+	// Respect the users PYTHONPATH if set
+	if currentValue, set := os.LookupEnv("PYTHONPATH"); set {
+		path = fmt.Sprintf("%s%c%s", currentValue, os.PathListSeparator, path)
+	}
+	cmd.Env = append(os.Environ(), fmt.Sprintf("PYTHONPATH=%s", path))
+}
diff --git a/pkg/scraper/script.go b/pkg/scraper/script.go
index 70bebfe4d..bfb03ee3a 100644
--- a/pkg/scraper/script.go
+++ b/pkg/scraper/script.go
@@ -44,6 +44,8 @@ func (s *scriptScraper) runScraperScript(ctx context.Context, inString string, o
 			logger.Warnf("%s", err)
 		} else {
 			cmd = p.Command(context.TODO(), command[1:])
+			envVariable, _ := filepath.Abs(filepath.Dir(filepath.Dir(s.config.path)))
+			python.AppendPythonPath(cmd, envVariable)
 		}
 	}