diff --git a/data/custom_extra_functionality.py b/data/custom_extra_functionality.py
index 0ac24cea..e7a9df0f 100644
--- a/data/custom_extra_functionality.py
+++ b/data/custom_extra_functionality.py
@@ -4,7 +4,7 @@
 # SPDX-License-Identifier: MIT
 
 
-def custom_extra_functionality(n, snapshots):
+def custom_extra_functionality(n, snapshots, snakemake):
     """
     Add custom extra functionality constraints.
     """
diff --git a/scripts/solve_network.py b/scripts/solve_network.py
index ecf56a24..4b988666 100644
--- a/scripts/solve_network.py
+++ b/scripts/solve_network.py
@@ -835,7 +835,7 @@ def extra_functionality(n, snapshots):
         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)
+        module.custom_extra_functionality(n, snapshots, snakemake)
 
 
 def solve_network(n, config, solving, opts="", **kwargs):