take the suggestions from the formatter

Author: yaoyi92

src/atomate2/forcefields/jobs.py

@@ -5,7 +5,7 @@
 import logging
 import warnings
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 from ase.io import Trajectory as AseTrajectory
 from ase.units import GPa as _GPa_to_eV_per_A3
@@ -15,7 +15,11 @@
 
 from atomate2.ase.jobs import AseRelaxMaker
 from atomate2.forcefields import MLFF, _get_formatted_ff_name
-from atomate2.forcefields.schemas import ForceFieldStructureTaskDocument, ForceFieldMoleculeTaskDocument, ForceFieldTaskDocument
+from atomate2.forcefields.schemas import (
+    ForceFieldStructureTaskDocument, 
+    ForceFieldMoleculeTaskDocument, 
+    ForceFieldTaskDocument
+)
 from atomate2.forcefields.utils import ase_calculator, revert_default_dtype
 
 if TYPE_CHECKING:
@@ -72,9 +76,7 @@ def make(structure):
     callable
         A decorated version of the make function that will generate forcefield jobs.
     """
-    return job(
-        method, data=_FORCEFIELD_DATA_OBJECTS
-    )
+    return job(method, data=_FORCEFIELD_DATA_OBJECTS)
 
 
 @dataclass
@@ -147,7 +149,7 @@ def __post_init__(self) -> None:
     @forcefield_job
     def make(
         self, structure: Structure, prev_dir: str | Path | None = None
-    ) -> Union[ForceFieldStructureTaskDocument, ForceFieldMoleculeTaskDocument]:
+    ) -> ForceFieldStructureTaskDocument | ForceFieldMoleculeTaskDocument:
         """
         Perform a relaxation of a structure using a force field.
 

src/atomate2/forcefields/md.py

@@ -4,7 +4,7 @@
 
 import warnings
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 from jobflow import job
 from monty.dev import deprecated
@@ -15,7 +15,11 @@
     _DEFAULT_CALCULATOR_KWARGS,
     _FORCEFIELD_DATA_OBJECTS,
 )
-from atomate2.forcefields.schemas import ForceFieldStructureTaskDocument, ForceFieldMoleculeTaskDocument, ForceFieldTaskDocument
+from atomate2.forcefields.schemas import (
+    ForceFieldStructureTaskDocument, 
+    ForceFieldMoleculeTaskDocument, 
+    ForceFieldTaskDocument,
+)
 from atomate2.forcefields.utils import ase_calculator, revert_default_dtype
 
 if TYPE_CHECKING:
@@ -131,7 +135,7 @@ def make(
         self,
         structure: Structure,
         prev_dir: str | Path | None = None,
-    ) -> Union[ForceFieldStructureTaskDocument, ForceFieldMoleculeTaskDocument]:
+    ) -> ForceFieldStructureTaskDocument | ForceFieldMoleculeTaskDocument:
         """
         Perform MD on a structure using forcefields and jobflow.
 

src/atomate2/forcefields/schemas.py

@@ -2,16 +2,22 @@
 
 from __future__ import annotations
 
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 from emmet.core.utils import ValueEnum
 from emmet.core.vasp.calculation import StoreTrajectoryOption
 from monty.dev import deprecated
 from pydantic import Field
 from pymatgen.core import Structure, Molecule
 
-from atomate2.ase.schemas import AseObject, AseResult, AseMoleculeTaskDoc, AseStructureTaskDoc, AseTaskDoc
-from atomate2.ase.schemas import _task_doc_translation_keys
+from atomate2.ase.schemas import (
+    AseObject, 
+    AseResult, 
+    AseMoleculeTaskDoc, 
+    AseStructureTaskDoc, 
+    AseTaskDoc,
+    _task_doc_translation_keys,
+)
 from atomate2.forcefields import MLFF
 
 
@@ -161,7 +167,7 @@ def from_ase_compatible_result(
         store_trajectory: StoreTrajectoryOption = StoreTrajectoryOption.NO,
         tags: list[str] | None = None,
         **task_document_kwargs,
-    ) -> Union[ForceFieldStructureTaskDocument, ForceFieldMoleculeTaskDocument]:
+    ) -> ForceFieldStructureTaskDocument | ForceFieldMoleculeTaskDocument:
         """Create an AseTaskDoc for a task that has ASE-compatible outputs.
 
         Parameters
@@ -232,7 +238,7 @@ def from_ase_compatible_result(
     @classmethod
     def from_ase_task_doc(
         cls, ase_task_doc: AseTaskDoc, **task_document_kwargs
-    ) -> Union[ForceFieldStructureTaskDocument, ForceFieldMoleculeTaskDocument]:
+    ) -> ForceFieldStructureTaskDocument | ForceFieldMoleculeTaskDocument:
         """Create an AseStructureTaskDoc for a task that has ASE-compatible outputs.
 
         Parameters
@@ -255,7 +261,9 @@ def from_ase_task_doc(
             meta_class = ForceFieldMoleculeTaskDocument
             k = "molecule"
             task_document_kwargs.update(molecule=ase_task_doc.mol_or_struct)
-        task_document_kwargs.update({k: ase_task_doc.mol_or_struct, f"meta_{k}": ase_task_doc.mol_or_struct})
+        task_document_kwargs.update(
+            {k: ase_task_doc.mol_or_struct, f"meta_{k}": ase_task_doc.mol_or_struct}
+        )
         return getattr(meta_class, f"from_{k}")(**task_document_kwargs)
 
     @property