Generalising fieldset.from_modulefile to also allow custom modulenames

Author: erikvansebille

parcels/fieldset.py

@@ -1065,27 +1065,30 @@ def from_xarray_dataset(cls, ds, variables, dimensions, mesh='spherical', allow_
         return cls(u, v, fields=fields)
 
     @classmethod
-    def from_modulefile(cls, filename, **kwargs):
+    def from_modulefile(cls, filename, modulename="create_fieldset", **kwargs):
         """Initialises FieldSet data from a file containing a python module file with a create_fieldset() function.
 
         Parameters
         ----------
-        filename: path to a python file containing at least a create_fieldset() function,
-            which returns a FieldSet object.
+        filename: path to a python file containing at least a function which returns a FieldSet object.
+        modulename: name of the function in the python file that returns a FieldSet object. Default is "create_fieldset".
         """
         # check if filename exists
         if not path.exists(filename):
             raise IOError(f"FieldSet module file {filename} does not exist")
 
         # Importing the source file directly (following https://docs.python.org/3/library/importlib.html#importing-a-source-file-directly)
-        spec = importlib.util.spec_from_file_location("create_fieldset", filename)
+        spec = importlib.util.spec_from_file_location(modulename, filename)
         fieldset_module = importlib.util.module_from_spec(spec)
-        sys.modules['create_fieldset'] = fieldset_module
+        sys.modules[modulename] = fieldset_module
         spec.loader.exec_module(fieldset_module)
 
-        if not hasattr(fieldset_module, 'create_fieldset'):
-            raise IOError(f"FieldSet module {filename} does not contain a `create_fieldset` function")
-        return fieldset_module.create_fieldset(**kwargs)
+        if not hasattr(fieldset_module, modulename):
+            raise IOError(f"{filename} does not contain a {modulename} function")
+        fieldset = getattr(fieldset_module, modulename)(**kwargs)
+        if not isinstance(fieldset, FieldSet):
+            raise IOError(f"Module {filename}.{modulename} does not return a FieldSet object")
+        return fieldset
 
     def get_fields(self):
         """Returns a list of all the :class:`parcels.field.Field` and :class:`parcels.field.VectorField`

tests/test_data/fieldset_nemo_error.py

@@ -15,3 +15,7 @@ def random_function_name():
     variables = {'U': 'U', 'V': 'V'}
     dimensions = {'lon': 'glamf', 'lat': 'gphif', 'time': 'time_counter'}
     return parcels.FieldSet.from_nemo(filenames, variables, dimensions)
+
+
+def none_returning_function():
+    return None

tests/test_fieldset.py

@@ -209,12 +209,14 @@ def test_fieldset_from_modulefile():
     fieldset = FieldSet.from_modulefile(data_path + 'fieldset_nemo.py', indices=indices)
     assert fieldset.U.grid.lon.shape[1] == 4
 
-
-def test_fieldset_from_modulefile_error():
-    data_path = path.join(path.dirname(__file__), 'test_data/')
     with pytest.raises(IOError):
         FieldSet.from_modulefile(data_path + 'fieldset_nemo_error.py')
 
+    FieldSet.from_modulefile(data_path + 'fieldset_nemo_error.py', modulename='random_function_name')
+
+    with pytest.raises(IOError):
+        FieldSet.from_modulefile(data_path + 'fieldset_nemo_error.py', modulename='none_returning_function')
+
 
 def test_field_from_netcdf_fieldtypes():
     data_path = path.join(path.dirname(__file__), 'test_data/')