add pydantic ConfigDicts to models using deprecated or incorrect allo…

…w extra behavior

emmet-core/emmet/core/openff/tasks.py

@@ -2,22 +2,23 @@
 
 from __future__ import annotations
 
+import zlib
 from dataclasses import dataclass
 from datetime import datetime
 from typing import Optional
-from typing_extensions import Annotated
-import zlib
 
+from monty.json import MSONable
 from pydantic import (
     BaseModel,
+    ConfigDict,
     Field,
-    PlainValidator,
     PlainSerializer,
+    PlainValidator,
     WithJsonSchema,
 )
-from monty.json import MSONable
-
 from pymatgen.core import Structure
+from typing_extensions import Annotated
+
 from emmet.core.vasp.task_valid import TaskState  # type: ignore[import-untyped]
 
 
@@ -56,7 +57,7 @@ class MoleculeSpec(MSONable):
     openff_mol: str  # a tk.Molecule object serialized with to_json
 
 
-class MDTaskDocument(BaseModel, extra="allow"):  # type: ignore[call-arg]
+class MDTaskDocument(BaseModel):  # type: ignore[call-arg]
     """Definition of the OpenMM task document."""
 
     tags: Optional[list[str]] = Field(
@@ -109,6 +110,8 @@ class MDTaskDocument(BaseModel, extra="allow"):  # type: ignore[call-arg]
         description="Timestamp for the most recent calculation for this task document",
     )
 
+    model_config = ConfigDict(extra="allow")
+
 
 class ClassicalMDTaskDocument(MDTaskDocument):
     """Definition of the OpenMM task document."""

emmet-core/emmet/core/openmm/tasks.py

@@ -5,20 +5,20 @@
 import io
 from pathlib import Path
 from typing import Optional, Union
-import pandas as pd  # type: ignore[import-untyped]
 
 import openmm
+import pandas as pd  # type: ignore[import-untyped]
 from openmm import XmlSerializer
 from openmm.app import Simulation
 from openmm.app.pdbfile import PDBFile
-from emmet.core.vasp.task_valid import TaskState  # type: ignore[import-untyped]
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
 
 from emmet.core.openff import MDTaskDocument  # type: ignore[import-untyped]
 from emmet.core.openff.tasks import CompressedStr  # type: ignore[import-untyped]
+from emmet.core.vasp.task_valid import TaskState  # type: ignore[import-untyped]
 
 
-class CalculationInput(BaseModel, extra="allow"):  # type: ignore[call-arg]
+class CalculationInput(BaseModel):  # type: ignore[call-arg]
     """OpenMM input settings for a job, these are the attributes of the OpenMMMaker."""
 
     n_steps: Optional[int] = Field(
@@ -100,6 +100,8 @@ class CalculationInput(BaseModel, extra="allow"):  # type: ignore[call-arg]
         description="Whether to embed the trajectory blob in CalculationOutput.",
     )
 
+    model_config = ConfigDict(extra="allow")
+
 
 class CalculationOutput(BaseModel):
     """OpenMM calculation output files and extracted data."""

emmet-core/emmet/core/vasp/calculation.py

@@ -8,7 +8,7 @@
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import numpy as np
-from pydantic import BaseModel, Extra, Field
+from pydantic import BaseModel, ConfigDict, Field
 from pymatgen.command_line.bader_caller import bader_analysis_from_path
 from pymatgen.command_line.chargemol_caller import ChargemolAnalysis
 from pymatgen.core.lattice import Lattice
@@ -301,7 +301,7 @@ class ElectronPhononDisplacedStructures(BaseModel):
     )
 
 
-class ElectronicStep(BaseModel, extra=Extra.allow):  # type: ignore
+class ElectronicStep(BaseModel):  # type: ignore
     """Document defining the information at each electronic step.
 
     Note, not all the information will be available at every step.
@@ -324,8 +324,10 @@ class ElectronicStep(BaseModel, extra=Extra.allow):  # type: ignore
     e_wo_entrp: Optional[float] = Field(None, description="The energy without entropy.")
     e_0_energy: Optional[float] = Field(None, description="The internal energy.")
 
+    model_config = ConfigDict(extra="allow")
 
-class IonicStep(BaseModel, extra=Extra.allow):  # type: ignore
+
+class IonicStep(BaseModel):  # type: ignore
     """Document defining the information at each ionic step."""
 
     e_fr_energy: Optional[float] = Field(None, description="The free energy.")
@@ -342,6 +344,8 @@ class IonicStep(BaseModel, extra=Extra.allow):  # type: ignore
         None, description="The structure at this step."
     )
 
+    model_config = ConfigDict(extra="allow")
+
 
 class CalculationOutput(BaseModel):
     """Document defining VASP calculation outputs."""
@@ -582,9 +586,11 @@ def from_vasp_outputs(
             frequency_dependent_dielectric=freq_dependent_diel,
             elph_displaced_structures=elph_structures,
             dos_properties=dosprop_dict,
-            ionic_steps=vasprun.ionic_steps
-            if store_trajectory == StoreTrajectoryOption.NO
-            else None,
+            ionic_steps=(
+                vasprun.ionic_steps
+                if store_trajectory == StoreTrajectoryOption.NO
+                else None
+            ),
             locpot=locpot_avg,
             outcar=outcar_dict,
             run_stats=RunStatistics.from_outcar(outcar) if outcar else None,