apply requested changes

aiidaplugins · Dec 20, 2023 · be2e8a0 · be2e8a0
1 parent 7a5881c
commit be2e8a0
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 16 deletions.
diff --git a/aiida_flexpart/calculations/flexpart_ifs.py b/aiida_flexpart/calculations/flexpart_ifs.py
@@ -14,7 +14,7 @@
 from ..utils import fill_in_template_file
 
 
-class FlexpartIfsCalculation(CalcJob):
+class FlexpartIfsCalculation(engine.CalcJob):
     """AiiDA calculation plugin wrapping the FLEXPART IFS executable."""
     @classmethod
     def define(cls, spec):
@@ -99,14 +99,14 @@ def prepare_for_submission(self, folder):
             meteo_string_list.append(f'{path}{os.sep}')
             meteo_string_list.append(f'{path}/AVAILABLE')
 
-        codeinfo = datastructures.CodeInfo()
+        codeinfo = common.CodeInfo()
         codeinfo.cmdline_params = meteo_string_list
         codeinfo.code_uuid = self.inputs.code.uuid
         codeinfo.stdout_name = self.metadata.options.output_filename
         codeinfo.withmpi = self.inputs.metadata.options.withmpi
 
         # Prepare a `CalcInfo` to be returned to the engine
-        calcinfo = datastructures.CalcInfo()
+        calcinfo = common.CalcInfo()
         calcinfo.codes_info = [codeinfo]
 
 

diff --git a/aiida_flexpart/parsers/flexpart_ifs.py b/aiida_flexpart/parsers/flexpart_ifs.py
@@ -4,12 +4,12 @@
 
 Register parsers via the "aiida.parsers" entry point in setup.json.
 """
-from aiida.engine, parsers, plugins, common, orm
+from aiida import parsers, plugins, common, orm, engine
 
-FlexpartCalculation = CalculationFactory('flexpart.ifs')
+FlexpartCalculation = plugins.CalculationFactory('flexpart.ifs')
 
 
-class FlexpartIfsParser(Parser):
+class FlexpartIfsParser(parsers.Parser):
     """
     Parser class for parsing output of calculation.
     """
@@ -24,7 +24,7 @@ def __init__(self, node):
         """
         super().__init__(node)
         if not issubclass(node.process_class, FlexpartCalculation):
-            raise exceptions.ParsingError('Can only parse FlexpartCalculation')
+            raise common.ParsingError('Can only parse FlexpartCalculation')
 
     def parse(self, **kwargs):
         """
@@ -39,21 +39,22 @@ def parse(self, **kwargs):
         files_expected = [output_filename]
         # Note: set(A) <= set(B) checks whether A is a subset of B
         if not set(files_expected) <= set(files_retrieved):
-            self.logger.error("Found files '{}', expected to find '{}'".format(
-                files_retrieved, files_expected))
+            self.logger.error(
+                f"Found files '{files_retrieved}', expected to find '{files_expected}'"
+            )
             return self.exit_codes.ERROR_MISSING_OUTPUT_FILES
 
         # check aiida.out content
         with self.retrieved.open(output_filename, 'r') as handle:
             content = handle.read()
-            output_node = SinglefileData(file=handle)
+            output_node = orm.SinglefileData(file=handle)
             if 'CONGRATULATIONS' not in content:
                 self.out('output_file', output_node)
-                return ExitCode(1)
+                return engine.ExitCode(1)
         # add output file
-        self.logger.info("Parsing '{}'".format(output_filename))
+        self.logger.info(f"Parsing '{output_filename}'")
         with self.retrieved.open(output_filename, 'rb') as handle:
-            output_node = SinglefileData(file=handle)
+            output_node = orm.SinglefileData(file=handle)
         self.out('output_file', output_node)
 
-        return ExitCode(0)
+        return engine.ExitCode(0)
diff --git a/examples/example_ifs.py b/examples/example_ifs.py
@@ -60,7 +60,7 @@ def test_run(flexpart_code):
         computer=flexpart_code.computer)
 
     # Set up calculation.
-    calc = CalculationFactory('flexpart.ifs')
+    calc = plugins.CalculationFactory('flexpart.ifs')
     builder = calc.get_builder()
     builder.code = flexpart_code
     builder.model_settings = {
@@ -85,7 +85,7 @@ def test_run(flexpart_code):
     builder.metadata.options.stash = {
         'source_list': ['aiida.out', 'grid_time_*.nc'],
         'target_base': f'/store/empa/em05/{user_name}/aiida_stash',
-        'stash_mode': StashMode.COPY.value,
+        'stash_mode': common.StashMode.COPY.value,
     }
     #builder.metadata.options.max_wallclock_seconds = 2000