Use typing to make sure component is Ocean

...in OceanIoStep and OceanModelStep

Author: xylar

polaris/ocean/model/ocean_io_step.py

@@ -7,6 +7,12 @@ class OceanIOStep(Step):
     A step that writes input and/or output files for Omega or MPAS-Ocean
     """
 
+    # make sure component is of type Ocean
+    component: Ocean
+
+    def __init__(self, component: Ocean, **kwargs):
+        super().__init__(component=component, **kwargs)
+
     def map_to_native_model_vars(self, ds):
         """
         If the model is Omega, rename dimensions and variables in a dataset
@@ -24,10 +30,6 @@ def map_to_native_model_vars(self, ds):
             The same dataset with variables renamed as appropriate for the
             ocean model being run
         """
-        if not isinstance(self.component, Ocean):
-            raise TypeError(
-                'component must be an instance of Ocean to map model vars'
-            )
         return self.component.map_to_native_model_vars(ds)
 
     def write_model_dataset(self, ds, filename):
@@ -43,10 +45,6 @@ def write_model_dataset(self, ds, filename):
         filename : str
             The path for the NetCDF file to write
         """
-        if not isinstance(self.component, Ocean):
-            raise TypeError(
-                'component must be an instance of Ocean to map model vars'
-            )
         self.component.write_model_dataset(ds, filename)
 
     def map_from_native_model_vars(self, ds):
@@ -65,10 +63,6 @@ def map_from_native_model_vars(self, ds):
         ds : xarray.Dataset
             The same dataset with variables named as expected in MPAS-Ocean
         """
-        if not isinstance(self.component, Ocean):
-            raise TypeError(
-                'component must be an instance of Ocean to map model vars'
-            )
         return self.component.map_from_native_model_vars(ds)
 
     def open_model_dataset(self, filename, **kwargs):
@@ -89,8 +83,4 @@ def open_model_dataset(self, filename, **kwargs):
         ds : xarray.Dataset
             The dataset with variables named as expected in MPAS-Ocean
         """
-        if not isinstance(self.component, Ocean):
-            raise TypeError(
-                'component must be an instance of Ocean to map model vars'
-            )
         return self.component.open_model_dataset(filename, **kwargs)

polaris/ocean/model/ocean_model_step.py

@@ -27,9 +27,12 @@ class OceanModelStep(ModelStep):
         working directory)
     """
 
+    # make sure component is of type Ocean
+    component: Ocean
+
     def __init__(
         self,
-        component,
+        component: Ocean,
         name,
         subdir=None,
         indir=None,
@@ -337,10 +340,6 @@ def validate_baselines(self):
         success : bool
             Whether the outputs were successfully validated against a baseline
         """
-        if not isinstance(self.component, Ocean):
-            raise TypeError(
-                'component must be an instance of Ocean to map model vars'
-            )
         # translate variable names to native model names
         validate_vars = {}
         for filename, vars in self.validate_vars.items():