Merge pull request #764 from dcs4cop/forman-xxx-let_server_scripts_im…

…port_modules

Allow scripts to import modules from script directory

CHANGES.md

@@ -2,6 +2,8 @@
 
 ### Fixes
 
+* xcube server Python scripts can now import modules from
+  the script's directory.
 * xcube serve correctly resolves relative paths to datasets (#758)
 
 ## Changes in 0.13.0.dev3

test/core/test_mldataset.py

@@ -1,11 +1,11 @@
 import os
-import threading
 import unittest
 
 import numpy as np
 import pandas as pd
 import xarray as xr
 
+from xcube.core.dsio import rimraf
 from xcube.core.gridmapping import GridMapping
 from xcube.core.mldataset import BaseMultiLevelDataset
 from xcube.core.mldataset import CombinedMultiLevelDataset
@@ -177,6 +177,48 @@ def input_ml_dataset_getter(ds_id):
         ml_ds1.close()
         ml_ds2.close()
 
+    def test_import(self):
+        script_dir = os.path.join(os.path.dirname(__file__), "test-code")
+
+        os.mkdir(script_dir)
+
+        with open(f"{script_dir}/code.py", "w") as fp:
+            fp.write(
+                "import module as m\n"
+                "\n"
+                "def compute_dataset(ds):\n"
+                "    return m.process_dataset(ds)\n"
+            )
+
+        with open(f"{script_dir}/module.py", "w") as fp:
+            fp.write(
+                "\n"
+                "def process_dataset(ds):\n"
+                "    return ds.copy()\n"
+            )
+
+        ds = _get_test_dataset()
+
+        ml_ds1 = BaseMultiLevelDataset(ds)
+
+        def input_ml_dataset_getter(ds_id):
+            if ds_id == "ml_ds1":
+                return ml_ds1
+            self.fail(f"unexpected ds_id={ds_id!r}")
+
+        try:
+            ml_ds2 = ComputedMultiLevelDataset(
+                f"{script_dir}/code.py",
+                "compute_dataset",
+                ["ml_ds1"],
+                input_ml_dataset_getter,
+                input_parameters=dict(),
+                ds_id="ml_ds2"
+            )
+            ml_ds2.get_dataset(0).compute()
+        finally:
+            rimraf(script_dir)
+
 
 def _get_test_dataset(var_names=('noise',)):
     w = 1440

xcube/core/mldataset.py

@@ -20,7 +20,9 @@
 # SOFTWARE.
 
 import math
+import os.path
 import threading
+import sys
 import uuid
 from abc import abstractmethod, ABCMeta
 from functools import cached_property
@@ -459,7 +461,11 @@ def _get_dataset_lazily(self,
 
 class ComputedMultiLevelDataset(LazyMultiLevelDataset):
     """
-    A multi-level dataset whose level datasets are a computed from a Python script.
+    A multi-level dataset whose level datasets are a computed from a Python
+    script.
+
+    The script can import other Python modules located in the same
+    directory as *script_path*.
     """
 
     def __init__(self,
@@ -481,6 +487,12 @@ def __init__(self,
             raise exception_type(
                 f"Failed to read Python code for in-memory dataset {ds_id!r} from {script_path!r}: {e}") from e
 
+        # Allow scripts to import modules from
+        # within directory
+        script_parent = os.path.dirname(script_path)
+        if script_parent not in sys.path:
+            sys.path = [script_parent] + sys.path
+
         global_env = dict()
         try:
             exec(python_code, global_env, None)