add mock_snakemake from pypsa-eur

scripts/helper.py

@@ -2,6 +2,62 @@
 import logging
 logger = logging.getLogger(__name__)
 
+import pandas as pd
+from pathlib import Path
+
+def mock_snakemake(rulename, **wildcards):
+    """
+    This function is expected to be executed from the 'scripts'-directory of '
+    the snakemake project. It returns a snakemake.script.Snakemake object,
+    based on the Snakefile.
+
+    If a rule has wildcards, you have to specify them in **wildcards.
+
+    Parameters
+    ----------
+    rulename: str
+        name of the rule for which the snakemake object should be generated
+    **wildcards:
+        keyword arguments fixing the wildcards. Only necessary if wildcards are
+        needed.
+    """
+    import snakemake as sm
+    import os
+    from pypsa.descriptors import Dict
+    from snakemake.script import Snakemake
+
+    script_dir = Path(__file__).parent.resolve()
+    assert Path.cwd().resolve() == script_dir, \
+      f'mock_snakemake has to be run from the repository scripts directory {script_dir}'
+    os.chdir(script_dir.parent)
+    for p in sm.SNAKEFILE_CHOICES:
+        if os.path.exists(p):
+            snakefile = p
+            break
+    workflow = sm.Workflow(snakefile)
+    workflow.include(snakefile)
+    workflow.global_resources = {}
+    rule = workflow.get_rule(rulename)
+    dag = sm.dag.DAG(workflow, rules=[rule])
+    wc = Dict(wildcards)
+    job = sm.jobs.Job(rule, dag, wc)
+
+    def make_accessable(*ios):
+        for io in ios:
+            for i in range(len(io)):
+                io[i] = os.path.abspath(io[i])
+
+    make_accessable(job.input, job.output, job.log)
+    snakemake = Snakemake(job.input, job.output, job.params, job.wildcards,
+                          job.threads, job.resources, job.log,
+                          job.dag.workflow.config, job.rule.name, None,)
+    # create log and output dir if not existent
+    for path in list(snakemake.log) + list(snakemake.output):
+        Path(path).parent.mkdir(parents=True, exist_ok=True)
+
+    os.chdir(script_dir)
+    return snakemake
+
 #https://stackoverflow.com/questions/20833344/fix-invalid-polygon-in-shapely
 #https://stackoverflow.com/questions/13062334/polygon-intersection-error-in-shapely-shapely-geos-topologicalerror-the-opera
 #https://shapely.readthedocs.io/en/latest/manual.html#object.buffer

scripts/solve_network.py

@@ -272,7 +272,7 @@ def solve_network(n, config=None, solver_log=None, opts=None):
                                                solver_name=solver_name,
                                                solver_logfile=solver_log,
                                                solver_options=solver_options,
-                                               solver_dir=tmpdir, 
+                                               solver_dir=tmpdir,
                                                extra_functionality=extra_functionality,
                                                formulation=solve_opts['formulation'])
                                                #extra_postprocessing=extra_postprocessing
@@ -371,19 +371,12 @@ def solve_network(n, config=None, solver_log=None, opts=None):
 if __name__ == "__main__":
     # Detect running outside of snakemake and mock snakemake for testing
     if 'snakemake' not in globals():
-        from vresutils.snakemake import MockSnakemake, Dict
-        snakemake = MockSnakemake(
-            wildcards=dict(network='elec', simpl='', clusters='39', lv='1.0',
-                           sector_opts='Co2L0-168H-T-H-B-I-solar3-dist1',
-                           co2_budget_name='b30b3', planning_horizons='2050'),
-            input=dict(network="pypsa-eur-sec/results/test/prenetworks_brownfield/elec_s{simpl}_{clusters}_lv{lv}__{sector_opts}_{co2_budget_name}_{planning_horizons}.nc"),
-            output=["results/networks/s{simpl}_{clusters}_lv{lv}_{sector_opts}_{co2_budget_name}_{planning_horizons}-test.nc"],
-            log=dict(gurobi="logs/elec_s{simpl}_{clusters}_lv{lv}_{sector_opts}_{co2_budget_name}_{planning_horizons}_gurobi-test.log",
-                     python="logs/elec_s{simpl}_{clusters}_lv{lv}_{sector_opts}_{co2_budget_name}_{planning_horizons}_python-test.log")
-        )
-        import yaml
-        with open('config.yaml', encoding='utf8') as f:
-            snakemake.config = yaml.safe_load(f)
+        from helpers import mock_snakemake
+        snakemake = mock_snakemake('solve_network',
+                                   network='elec', simpl='', clusters='37',
+                                   lv='1.0', opts='', planning_horizons='2020',
+                                   sector_opts='168H-T-H-B-I')
+
     tmpdir = snakemake.config['solving'].get('tmpdir')
     if tmpdir is not None:
         patch_pyomo_tmpdir(tmpdir)