From d145758fb7ff4bf126ddada8eec6d8f942c93f4f Mon Sep 17 00:00:00 2001
From: Fabian Neumann <fabian.neumann@outlook.de>
Date: Thu, 4 Jan 2024 09:00:31 +0100
Subject: [PATCH] gracefully handle absent extra_functionality file; add file
 to path

---
 rules/common.smk            | 7 +++++++
 rules/solve_electricity.smk | 4 +---
 rules/solve_myopic.smk      | 4 +---
 rules/solve_overnight.smk   | 5 +----
 rules/solve_perfect.smk     | 4 +---
 scripts/solve_network.py    | 2 ++
 6 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/rules/common.smk b/rules/common.smk
index 2c8cf69c..44e3a807 100644
--- a/rules/common.smk
+++ b/rules/common.smk
@@ -28,6 +28,13 @@ def memory(w):
         return int(factor * (10000 + 195 * int(w.clusters)))
 
 
+def input_custom_extra_functionality(w):
+    path = config["solving"]["options"].get("custom_extra_functionality", False)
+    if path:
+        return workflow.source_path(path)
+    return []
+
+
 # Check if the workflow has access to the internet by trying to access the HEAD of specified url
 def has_internet_access(url="www.zenodo.org") -> bool:
     import http.client as http_client
diff --git a/rules/solve_electricity.smk b/rules/solve_electricity.smk
index 2c956097..7f6092be 100644
--- a/rules/solve_electricity.smk
+++ b/rules/solve_electricity.smk
@@ -11,9 +11,7 @@ rule solve_network:
         co2_sequestration_potential=config["sector"].get(
             "co2_sequestration_potential", 200
         ),
-        custom_extra_functionality=workflow.source_path(
-            config["solving"]["options"].get("custom_extra_functionality", "")
-        ),
+        custom_extra_functionality=input_custom_extra_functionality,
     input:
         network=RESOURCES + "networks/elec_s{simpl}_{clusters}_ec_l{ll}_{opts}.nc",
         config=RESULTS + "config.yaml",
diff --git a/rules/solve_myopic.smk b/rules/solve_myopic.smk
index afa8ad2c..7ca8857d 100644
--- a/rules/solve_myopic.smk
+++ b/rules/solve_myopic.smk
@@ -88,9 +88,7 @@ rule solve_sector_network_myopic:
         co2_sequestration_potential=config["sector"].get(
             "co2_sequestration_potential", 200
         ),
-        custom_extra_functionality=workflow.source_path(
-            config["solving"]["options"].get("custom_extra_functionality", "")
-        ),
+        custom_extra_functionality=input_custom_extra_functionality,
     input:
         network=RESULTS
         + "prenetworks-brownfield/elec_s{simpl}_{clusters}_l{ll}_{opts}_{sector_opts}_{planning_horizons}.nc",
diff --git a/rules/solve_overnight.smk b/rules/solve_overnight.smk
index fc2f74df..8686b205 100644
--- a/rules/solve_overnight.smk
+++ b/rules/solve_overnight.smk
@@ -2,7 +2,6 @@
 #
 # SPDX-License-Identifier: MIT
 
-
 rule solve_sector_network:
     params:
         solving=config["solving"],
@@ -11,9 +10,7 @@ rule solve_sector_network:
         co2_sequestration_potential=config["sector"].get(
             "co2_sequestration_potential", 200
         ),
-        custom_extra_functionality=workflow.source_path(
-            config["solving"]["options"].get("custom_extra_functionality", "")
-        ),
+        custom_extra_functionality=input_custom_extra_functionality,
     input:
         network=RESULTS
         + "prenetworks/elec_s{simpl}_{clusters}_l{ll}_{opts}_{sector_opts}_{planning_horizons}.nc",
diff --git a/rules/solve_perfect.smk b/rules/solve_perfect.smk
index 63be5cc1..a7856fa9 100644
--- a/rules/solve_perfect.smk
+++ b/rules/solve_perfect.smk
@@ -118,9 +118,7 @@ rule solve_sector_network_perfect:
         co2_sequestration_potential=config["sector"].get(
             "co2_sequestration_potential", 200
         ),
-        custom_extra_functionality=workflow.source_path(
-            config["solving"]["options"].get("custom_extra_functionality", "")
-        ),
+        custom_extra_functionality=input_custom_extra_functionality,
     input:
         network=RESULTS
         + "prenetworks-brownfield/elec_s{simpl}_{clusters}_l{ll}_{opts}_{sector_opts}_brownfield_all_years.nc",
diff --git a/scripts/solve_network.py b/scripts/solve_network.py
index 5a045577..2f170dff 100644
--- a/scripts/solve_network.py
+++ b/scripts/solve_network.py
@@ -30,6 +30,7 @@ import importlib
 import logging
 import os
 import re
+import sys
 
 import numpy as np
 import pandas as pd
@@ -831,6 +832,7 @@ def extra_functionality(n, snapshots):
     if snakemake.params.custom_extra_functionality:
         source_path = snakemake.params.custom_extra_functionality
         assert os.path.exists(source_path), f"{source_path} does not exist"
+        sys.path.append(os.path.dirname(source_path))
         module_name = os.path.splitext(os.path.basename(source_path))[0]
         module = importlib.import_module(module_name)
         module.custom_extra_functionality(n, snapshots)