Fix failing tests from Pydantic v2 migration

src/atomate2/cp2k/schemas/calculation.py

@@ -8,7 +8,7 @@
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 from jobflow.utils import ValueEnum
-from pydantic import BaseModel, Field, validator
+from pydantic import BaseModel, Field, field_validator
 from pymatgen.command_line.bader_caller import BaderAnalysis
 from pymatgen.core.structure import Molecule, Structure
 from pymatgen.core.trajectory import Trajectory
@@ -82,16 +82,16 @@ class CalculationInput(BaseModel):
         description="CP2K global parameters used in the last calc of this task.",
     )
 
-    @validator("atomic_kind_info")
-    def remove_unnecessary(self, atomic_kind_info):
+    @field_validator("atomic_kind_info")
+    def remove_unnecessary(cls, atomic_kind_info):
         """Remove unnecessary entry from atomic_kind_info."""
         for k in atomic_kind_info:
             if "total_pseudopotential_energy" in atomic_kind_info[k]:
                 del atomic_kind_info[k]["total_pseudopotential_energy"]
         return atomic_kind_info
 
-    @validator("dft")
-    def cleanup_dft(self, dft):
+    @field_validator("dft")
+    def cleanup_dft(cls, dft):
         """Convert UKS strings to UKS=True."""
         if any(v.upper() == "UKS" for v in dft.values()):
             dft["UKS"] = True

src/atomate2/cp2k/schemas/task.py

@@ -234,51 +234,59 @@ def from_cp2k_calc_doc(cls, calc_doc: Calculation) -> "OutputSummary":
 class TaskDocument(StructureMetadata, MoleculeMetadata):
     """Definition of CP2K task document."""
 
-    dir_name: str = Field(None, description="The directory for this CP2K task")
+    dir_name: Optional[str] = Field(
+        None, description="The directory for this CP2K task"
+    )
     last_updated: str = Field(
         default_factory=datetime_str,
         description="Timestamp for this task document was last updated",
     )
-    completed_at: str = Field(
+    completed_at: Optional[str] = Field(
         None, description="Timestamp for when this task was completed"
     )
-    input: InputSummary = Field(None, description="The input to the first calculation")
-    output: OutputSummary = Field(
+    input: Optional[InputSummary] = Field(
+        None, description="The input to the first calculation"
+    )
+    output: Optional[OutputSummary] = Field(
         None, description="The output of the final calculation"
     )
     structure: Union[Structure, Molecule] = Field(
         None, description="Final output structure from the task"
     )
-    state: Status = Field(None, description="State of this task")
-    included_objects: List[Cp2kObject] = Field(
+    state: Optional[Status] = Field(None, description="State of this task")
+    included_objects: Optional[List[Cp2kObject]] = Field(
         None, description="List of CP2K objects included with this task document"
     )
-    cp2k_objects: Dict[Cp2kObject, Any] = Field(
+    cp2k_objects: Optional[Dict[Cp2kObject, Any]] = Field(
         None, description="CP2K objects associated with this task"
     )
-    entry: ComputedEntry = Field(
+    entry: Optional[ComputedEntry] = Field(
         None, description="The ComputedEntry from the task doc"
     )
-    analysis: AnalysisSummary = Field(
+    analysis: Optional[AnalysisSummary] = Field(
         None, description="Summary of structural relaxation and forces"
     )
-    run_stats: Dict[str, RunStatistics] = Field(
+    run_stats: Optional[Dict[str, RunStatistics]] = Field(
         None,
         description="Summary of runtime statistics for each calculation in this task",
     )
-    orig_inputs: Dict[str, Cp2kInput] = Field(
+    orig_inputs: Optional[Dict[str, Cp2kInput]] = Field(
         None, description="Summary of the original CP2K inputs written by custodian"
     )
-    task_label: str = Field(None, description="A description of the task")
-    tags: List[str] = Field(None, description="Metadata tags for this task document")
-    author: str = Field(None, description="Author extracted from transformations")
-    icsd_id: str = Field(
+    task_label: Optional[str] = Field(None, description="A description of the task")
+    tags: Optional[List[str]] = Field(
+        None, description="Metadata tags for this task document"
+    )
+    author: Optional[str] = Field(
+        None, description="Author extracted from transformations"
+    )
+    icsd_id: Optional[str] = Field(
         None, description="International crystal structure database id of the structure"
     )
-    calcs_reversed: List[Calculation] = Field(
+    calcs_reversed: Optional[List[Calculation]] = Field(
         None, description="The inputs and outputs for all CP2K runs in this task."
     )
-    transformations: Dict[str, Any] = Field(
+    transformations: Optional[Dict[str, Any]] = Field(
         None,
         description="Information on the structural transformations, parsed from a "
         "transformations.json file",
@@ -288,7 +296,7 @@ class TaskDocument(StructureMetadata, MoleculeMetadata):
         description="Information on the custodian settings used to run this "
         "calculation, parsed from a custodian.json file",
     )
-    additional_json: Dict[str, Any] = Field(
+    additional_json: Optional[Dict[str, Any]] = Field(
         None, description="Additional json loaded from the calculation directory"
     )
     schema: str = Field(