Support pydantic 2

Phys/DaVinci/python/DaVinci/LbExec.py

@@ -10,7 +10,7 @@
 ###############################################################################
 from contextlib import contextmanager
 from typing import Optional, Any
-from pydantic import root_validator, PrivateAttr  # type: ignore[import]
+from pydantic import model_validator, PrivateAttr  # type: ignore[import]
 import logging
 from GaudiConf.LbExec import Options as DefaultOptions, TestOptionsBase  # type: ignore[import]
 from GaudiConf.LbExec import FileFormats, InputProcessTypes  # type: ignore[import]
@@ -53,26 +53,27 @@ class Options(DefaultOptions):
     metainfo_additional_tags: Optional[list] = []
     _input_config: Any = PrivateAttr()
 
-    @root_validator(pre=False, skip_on_failure=True)
-    def _stream_default(cls, values):
+    @model_validator(mode="before")
+    @classmethod
+    def _stream_default(cls, data):
         """
         This is a validator that sets the default "stream" value based on "input_process"
 
         Args:
-          values (dict): User-specified attributes of the Options object.
+          data (dict): User-specified attributes of the Options object.
 
         Returns:
           dict: Modified attributes of the Options object.
         """
-        input_process = values.get("input_process")
-        input_stream = values.get("input_stream")
+        input_process = data.get("input_process")
+        input_stream = data.get("input_stream")
 
-        if not input_stream.islower():
-            logging.getLogger(__name__).warning(
-                f"input_stream is set to '{input_stream}', input_stream should be lower case. It will be made lower case for you in the next step. If you think this is wrong contact DaVinci maintainers."
-            )
-
-        values["input_stream"] = input_stream.lower()
+        if input_stream:
+            if not input_stream.islower():
+                logging.getLogger(__name__).warning(
+                    f"input_stream is set to '{input_stream}', input_stream should be lower case. It will be made lower case for you in the next step. If you think this is wrong contact DaVinci maintainers."
+                )
+            data["input_stream"] = input_stream.lower()
 
         if (
             input_process
@@ -86,81 +87,62 @@ class Options(DefaultOptions):
             logging.getLogger(__name__).warning(
                 f"input_stream is set to '{input_stream}', but will be reset to '' because current input_process = {input_process}"
             )
-            values["input_stream"] = ""
+            data["input_stream"] = ""
 
-        return values
+        return data
 
-    @root_validator(pre=False, skip_on_failure=True)
-    def validate_input_output(cls, values):
+    @model_validator(mode="after")
+    def validate_input_output(self):
         """
         Validator for the consistency of the input process
         and input/output file formats for Sprucing jobs.
-
-        Args:
-          values (dict): User-specified attributes of the Options object.
-
-        Returns:
-          dict: Modified attributes of the Options object.
         """
-        input_process = values.get("input_process", None)
-        input_type = values.get("input_type", FileFormats.NONE)
-        output_type = values.get("output_type", FileFormats.ROOT)
-        simulation = values.get("simulation", None)
-
         # Jobs on Spruced samples are DST -> DST only
-        if input_process in {InputProcessTypes.TurboPass, InputProcessTypes.Spruce}:
-            if input_type != FileFormats.ROOT:
+        if self.input_process in {
+            InputProcessTypes.TurboPass,
+            InputProcessTypes.Spruce,
+        }:
+            if self.input_type != FileFormats.ROOT:
                 raise ValueError(
-                    f"Jobs on Spruced data are DST -> ROOT but input_type={input_type!r} was given."
+                    f"Jobs on Spruced data are DST -> ROOT but input_type={self.input_type!r} was given."
                 )
-            if output_type != FileFormats.ROOT:
+            if self.output_type != FileFormats.ROOT:
                 raise ValueError(
-                    f"Jobs on Spruced data are DST -> ROOT but output_type={output_type!r} was given."
+                    f"Jobs on Spruced data are DST -> ROOT but output_type={self.output_type!r} was given."
                 )
         if (
-            input_process
+            self.input_process
             in [InputProcessTypes.TurboSpruce, InputProcessTypes.TurboPass]
-            and simulation
-        ) and not issubclass(cls, TestOptions):
+            and self.simulation
+        ) and not issubclass(type(self), TestOptions):
             raise ValueError(
-                f"TurboSpruce and TurboPass is not supported for simulation. input_process={input_process!r} and simulation={simulation!r} was given."
+                f"TurboSpruce and TurboPass is not supported for simulation. input_process={self.input_process!r} and simulation={self.simulation!r} was given."
             )
 
-        return values
+        return self
 
-    @root_validator(pre=False, skip_on_failure=True)
-    def validate_fsr_options(cls, values):
+    @model_validator(mode="after")
+    def validate_fsr_options(self):
         """
         Validator for the consistency of writing FSRs and other settings.
-
-        Args:
-          values (dict): User-specified attributes of the Options object.
-
-        Returns:
-          dict: Modified attributes of the Options object.
         """
-        write_fsr = values.get("write_fsr", False)
-        merge_genfsr = values.get("merge_genfsr", False)
-        output_file = values.get("output_file", None)
-        simulation = values.get("simulation", False)
-
-        if write_fsr:
-            if not output_file:  # Catch the default None but also ""
+        if self.write_fsr:
+            if not self.output_file:  # Catch the default None but also ""
                 raise ValueError(
                     "FSRs requested to be written but no output file was provided."
                 )
 
-        if merge_genfsr:
-            if not simulation:
+        if self.merge_genfsr:
+            if not self.simulation:
                 raise ValueError(
                     "Merging of Generator FSRs is only relevant/valid for simulation."
                 )
-            if not output_file:  # Catch the default None but also ""
+            if not self.simulation:  # Catch the default None but also ""
                 raise ValueError(
                     "Merging of Generator FSRs requested to be written but no output file was provided."
                 )
 
-        return values
+        return self
 
     @contextmanager
     def apply_binds(self):
@@ -204,7 +186,7 @@ class Options(DefaultOptions):
             yield
 
 
-class TestOptions(Options, TestOptionsBase):
+class TestOptions(TestOptionsBase, Options):
     """Specialized Options class for Davinci Tests"""
 
     pass