switch to pydantic v2

Author: yaugenst-flex

tests/config/test_loader.py

@@ -2,8 +2,8 @@
 
 from pathlib import Path
 
-import pydantic.v1 as pd
 from click.testing import CliRunner
+from pydantic import Field
 
 from tidy3d.config import get_manager, reload_config
 from tidy3d.config import registry as config_registry
@@ -49,7 +49,7 @@ def test_profile_preserves_comments(config_manager, mock_config_dir):
     class ProfileComment(ConfigSection):
         """Profile comment plugin."""
 
-        knob: int = pd.Field(
+        knob: int = Field(
             1,
             description="Profile knob description.",
             json_schema_extra={"persist": True},
@@ -83,7 +83,7 @@ def test_cli_reset_config(mock_config_dir):
     class CLIPlugin(ConfigSection):
         """CLI plugin configuration."""
 
-        knob: int = pd.Field(
+        knob: int = Field(
             3,
             description="CLI knob description.",
             json_schema_extra={"persist": True},
@@ -120,7 +120,7 @@ def test_plugin_descriptions(mock_config_dir):
     class CommentPlugin(ConfigSection):
         """Comment plugin configuration."""
 
-        knob: int = pd.Field(
+        knob: int = Field(
             3,
             description="Plugin knob description.",
             json_schema_extra={"persist": True},

tests/config/test_plugins.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
-import pydantic.v1 as pd
 import toml
+from pydantic import Field
 
 from tidy3d.config.__init__ import get_manager, reload_config
 from tidy3d.config.registry import get_sections, register_plugin
@@ -14,8 +14,8 @@ def ensure_dummy_plugin():
 
     @register_plugin("dummy")
     class DummyPlugin(ConfigSection):
-        enabled: bool = pd.Field(False, json_schema_extra={"persist": True})
-        precision: int = pd.Field(1, json_schema_extra={"persist": True})
+        enabled: bool = Field(False, json_schema_extra={"persist": True})
+        precision: int = Field(1, json_schema_extra={"persist": True})
 
 
 def test_plugin_defaults_available(mock_config_dir):

tests/config/test_profiles.py

@@ -20,4 +20,6 @@ def test_save_custom_profile(config_manager):
     config_manager.save()
 
     profile_path = config_manager.config_dir / "profiles" / "customer.toml"
-    assert not profile_path.exists()
+    assert profile_path.exists()
+    data = toml.load(profile_path)
+    assert data["logging"]["level"] == "DEBUG"

tidy3d/config/manager.py

@@ -8,9 +8,9 @@
 from collections.abc import Iterable
 from copy import deepcopy
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any, Optional, get_args, get_origin
 
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 
 from tidy3d.log import log
 
@@ -61,7 +61,7 @@ def dict(self, *args, **kwargs):  # type: ignore[override]
         model = self._manager._get_model(self._path)
         if model is None:
             return {}
-        return model.dict(*args, **kwargs)
+        return model.model_dump(*args, **kwargs)
 
 
 class PluginsAccessor:
@@ -380,7 +380,7 @@ def __setattr__(self, name: str, value: Any) -> None:
             return
         if name in self._section_models:
             if isinstance(value, BaseModel):
-                payload = value.dict(exclude_unset=False)
+                payload = value.model_dump(exclude_unset=False)
             else:
                 payload = value
             self.update_section(name, **payload)
@@ -404,16 +404,33 @@ def _deep_get(tree: dict[str, Any], path: Iterable[str]) -> Optional[dict[str, A
     return node if isinstance(node, dict) else None
 
 
+def _resolve_model_type(annotation: Any) -> Optional[type[BaseModel]]:
+    """Return the first BaseModel subclass found in an annotation (if any)."""
+
+    if isinstance(annotation, type) and issubclass(annotation, BaseModel):
+        return annotation
+
+    origin = get_origin(annotation)
+    if origin is None:
+        return None
+
+    for arg in get_args(annotation):
+        nested = _resolve_model_type(arg)
+        if nested is not None:
+            return nested
+    return None
+
+
 def _serialize_value(value: Any) -> Any:
     if isinstance(value, BaseModel):
-        return value.dict(exclude_unset=False)
+        return value.model_dump(exclude_unset=False)
     if hasattr(value, "get_secret_value"):
         return value.get_secret_value()
     return value
 
 
 def _model_dict(model: BaseModel) -> dict[str, Any]:
-    data = model.dict(exclude_unset=False)
+    data = model.model_dump(exclude_unset=False)
     for key, value in list(data.items()):
         if hasattr(value, "get_secret_value"):
             data[key] = value.get_secret_value()
@@ -422,10 +439,10 @@ def _model_dict(model: BaseModel) -> dict[str, Any]:
 
 def _extract_persisted(schema: type[BaseModel], data: dict[str, Any]) -> dict[str, Any]:
     persisted: dict[str, Any] = {}
-    for field_name, field in schema.__fields__.items():
-        extra = getattr(field.field_info, "extra", {}) or {}
-        schema_extra = extra.get("json_schema_extra", {})
-        persist = schema_extra.get("persist") if isinstance(schema_extra, dict) else False
+    for field_name, field in schema.model_fields.items():
+        schema_extra = field.json_schema_extra or {}
+        annotation = field.annotation
+        persist = bool(schema_extra.get("persist")) if isinstance(schema_extra, dict) else False
         if not persist:
             continue
         if field_name not in data:
@@ -435,10 +452,10 @@ def _extract_persisted(schema: type[BaseModel], data: dict[str, Any]) -> dict[st
             persisted[field_name] = None
             continue
 
-        field_type = getattr(field, "type_", None)
-        if isinstance(field_type, type) and issubclass(field_type, BaseModel):
+        nested_type = _resolve_model_type(annotation)
+        if nested_type is not None:
             nested_source = value if isinstance(value, dict) else {}
-            nested_persisted = _extract_persisted(field_type, nested_source)
+            nested_persisted = _extract_persisted(nested_type, nested_source)
             if nested_persisted:
                 persisted[field_name] = nested_persisted
             continue

tidy3d/config/registry.py

@@ -4,7 +4,7 @@
 
 from typing import Callable, Optional, TypeVar
 
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 
 T = TypeVar("T", bound=BaseModel)