switch to pydantic v2 (#485)

Closes #483.

Author: skshetry

gto/_pydantic.py

@@ -1,6 +1,7 @@
 from datetime import datetime
 from typing import Any, Dict, FrozenSet, List, Optional, Sequence, Union
 
+from pydantic import BaseModel, ConfigDict
 from scmrepo.git import Git
 
 from gto.config import RegistryConfig
@@ -12,7 +13,6 @@
 )
 from gto.versions import SemVer
 
-from ._pydantic import BaseModel
 from .exceptions import (
     ArtifactNotFound,
     ManyVersions,
@@ -41,7 +41,7 @@ def event(self):
         return self.__class__.__name__.lower()
 
     def dict_state(self, exclude=None):
-        state = self.dict(exclude=exclude)
+        state = self.model_dump(exclude=exclude)
         state["event"] = self.event
         return state
 
@@ -178,7 +178,7 @@ def ref(self):
         return self.authoring_event.ref
 
     def dict_state(self, exclude=None):
-        version = self.dict(exclude=exclude)
+        version = self.model_dump(exclude=exclude)
         version["is_active"] = self.is_active
         version["activated_at"] = self.activated_at
         version["created_at"] = self.created_at
@@ -565,9 +565,7 @@ def find_version_at_commit(
 
 class BaseRegistryState(BaseModel):
     artifacts: Dict[str, Artifact] = {}
-
-    class Config:
-        arbitrary_types_allowed = True
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
     def add_artifact(self, name):
         self.artifacts[name] = Artifact(artifact=name, versions=[])
@@ -623,9 +621,7 @@ class BaseManager(BaseModel):
     scm: Git
     actions: FrozenSet[Action]
     config: RegistryConfig
-
-    class Config:
-        arbitrary_types_allowed = True
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
     def update_state(self, state: BaseRegistryState) -> BaseRegistryState:
         raise NotImplementedError

gto/base.py

@@ -810,7 +810,7 @@ def stages(
 def print_state(repo: str = option_repo):
     """Technical cmd: Print current registry state."""
     state = make_ready_to_serialize(
-        gto.api._get_state(repo).dict()  # pylint: disable=protected-access
+        gto.api._get_state(repo).model_dump()  # pylint: disable=protected-access
     )
     format_echo(state, "json")
 
@@ -833,7 +833,7 @@ def doctor(
             echo(f"{EMOJI_FAIL} Fail to parse config")
         echo("---------------------------------")
 
-    gto.api._get_state(repo).dict()  # pylint: disable=protected-access
+    gto.api._get_state(repo).model_dump()  # pylint: disable=protected-access
     with cli_echo():
         echo(f"{EMOJI_OK} No issues found")
 

gto/cli.py

@@ -1,16 +1,22 @@
 # pylint: disable=no-self-argument, inconsistent-return-statements, invalid-name, import-outside-toplevel
 import pathlib
-from pathlib import Path
 from typing import Any, Dict, List, Optional
 
+from pydantic import BaseModel, Field, field_validator
+from pydantic_settings import (
+    BaseSettings,
+    PydanticBaseSettingsSource,
+    SettingsConfigDict,
+)
+from pydantic_settings import (
+    YamlConfigSettingsSource as _YamlConfigSettingsSource,
+)
 from ruamel.yaml import YAML
 
 from gto.constants import assert_name_is_valid
 from gto.exceptions import UnknownStage, UnknownType, WrongConfig
 from gto.ext import EnrichmentReader, find_enrichment_types, find_enrichments
 
-from ._pydantic import BaseModel, BaseSettings, InitSettingsSource, validator
-
 yaml = YAML(typ="safe", pure=True)
 yaml.default_flow_style = False
 
@@ -27,45 +33,47 @@ def load(self) -> EnrichmentReader:
 
 class NoFileConfig(BaseSettings):  # type: ignore[valid-type]
     INDEX: str = "artifacts.yaml"
-    TYPES: Optional[List[str]] = None
-    STAGES: Optional[List[str]] = None
+    CONFIG_FILE_NAME: Optional[str] = CONFIG_FILE_NAME
     LOG_LEVEL: str = "INFO"
     DEBUG: bool = False
-    ENRICHMENTS: List[EnrichmentConfig] = []
-    AUTOLOAD_ENRICHMENTS: bool = True
-    CONFIG_FILE_NAME: Optional[str] = CONFIG_FILE_NAME
     EMOJIS: bool = True
 
-    class Config:
-        env_prefix = "gto_"
+    types: Optional[List[str]] = None
+    stages: Optional[List[str]] = None
+    enrichments: List[EnrichmentConfig] = Field(default_factory=list)
+    autoload_enrichments: bool = True
+
+    model_config = SettingsConfigDict(env_prefix="gto_")
 
     def assert_type(self, name):
         assert_name_is_valid(name)
         # pylint: disable-next=unsupported-membership-test
-        if self.TYPES is not None and name not in self.TYPES:
-            raise UnknownType(name, self.TYPES)
+        if self.types is not None and name not in self.types:
+            raise UnknownType(name, self.types)
 
     def assert_stage(self, name):
         assert_name_is_valid(name)
         # pylint: disable-next=unsupported-membership-test
-        if self.STAGES is not None and name not in self.STAGES:
-            raise UnknownStage(name, self.STAGES)
+        if self.stages is not None and name not in self.stages:
+            raise UnknownStage(name, self.stages)
 
     @property
-    def enrichments(self) -> Dict[str, EnrichmentReader]:
-        res = {e.source: e for e in (e.load() for e in self.ENRICHMENTS)}
-        if self.AUTOLOAD_ENRICHMENTS:
+    def enrichments_(self) -> Dict[str, EnrichmentReader]:
+        res = {e.source: e for e in (e.load() for e in self.enrichments)}
+        if self.autoload_enrichments:
             return {**find_enrichments(), **res}
         return res
 
-    @validator("TYPES")
+    @field_validator("types")
+    @classmethod
     def types_are_valid(cls, v):  # pylint: disable=no-self-use
         if v:
             for name in v:
                 assert_name_is_valid(name)
         return v
 
-    @validator("STAGES")
+    @field_validator("stages")
+    @classmethod
     def stages_are_valid(cls, v):  # pylint: disable=no-self-use
         if v:
             for name in v:
@@ -77,61 +85,48 @@ def check_index_exist(self, repo: str):
         return index.exists() and index.is_file()
 
 
-def _set_location_init_source(init_source: InitSettingsSource):
-    def inner(settings: "RegistryConfig"):
-        if "CONFIG_FILE_NAME" in init_source.init_kwargs:
-            settings.__dict__["CONFIG_FILE_NAME"] = init_source.init_kwargs[
-                "CONFIG_FILE_NAME"
-            ]
-        return {}
-
-    return inner
+class YamlConfigSettingsSource(_YamlConfigSettingsSource):
+    def _read_file(self, file_path: pathlib.Path) -> dict[str, Any]:
+        with open(file_path, encoding=self.yaml_file_encoding) as yaml_file:
+            return yaml.load(yaml_file) or {}
 
 
-def config_settings_source(settings: "RegistryConfig") -> Dict[str, Any]:
-    """
-    A simple settings source that loads variables from a yaml file in GTO DIR
-    """
-
-    encoding = settings.__config__.env_file_encoding
-    config_file = getattr(settings, "CONFIG_FILE_NAME", CONFIG_FILE_NAME)
-    if not isinstance(config_file, Path):
-        config_file = Path(config_file)
-    if not config_file.exists():
-        return {}
-    conf = yaml.load(config_file.read_text(encoding=encoding))
-
-    return {k.upper(): v for k, v in conf.items()} if conf else {}
+class RegistryConfig(NoFileConfig):
+    model_config = SettingsConfigDict(env_prefix="gto_", env_file_encoding="utf-8")
 
+    def config_file_exists(self):
+        config = pathlib.Path(self.CONFIG_FILE_NAME)
+        return config.exists() and config.is_file()
 
-class RegistryConfig(NoFileConfig):
-    class Config:
-        env_prefix = "gto_"
-        env_file_encoding = "utf-8"
 
+def read_registry_config(config_file_name) -> "RegistryConfig":
+    class _RegistryConfig(RegistryConfig):
         @classmethod
-        def customise_sources(
+        def settings_customise_sources(
             cls,
-            init_settings,
-            env_settings,
-            file_secret_settings,
+            settings_cls: type[BaseSettings],
+            init_settings: PydanticBaseSettingsSource,
+            env_settings: PydanticBaseSettingsSource,
+            dotenv_settings: PydanticBaseSettingsSource,
+            file_secret_settings: PydanticBaseSettingsSource,
         ):
+            encoding = getattr(settings_cls.model_config, "env_file_encoding", "utf-8")
             return (
-                _set_location_init_source(init_settings),
                 init_settings,
                 env_settings,
-                config_settings_source,
+                (
+                    YamlConfigSettingsSource(
+                        settings_cls,
+                        yaml_file=config_file_name,
+                        yaml_file_encoding=encoding,
+                    )
+                ),
+                dotenv_settings,
                 file_secret_settings,
             )
 
-    def config_file_exists(self):
-        config = pathlib.Path(self.CONFIG_FILE_NAME)
-        return config.exists() and config.is_file()
-
-
-def read_registry_config(config_file_name):
     try:
-        return RegistryConfig(CONFIG_FILE_NAME=config_file_name)
+        return _RegistryConfig(CONFIG_FILE_NAME=config_file_name)
     except Exception as e:  # pylint: disable=bare-except
         raise WrongConfig(config_file_name) from e
 

gto/config.py

@@ -2,9 +2,9 @@
 from enum import Enum
 from typing import Optional
 
-from gto.exceptions import ValidationError
+from pydantic import BaseModel
 
-from ._pydantic import BaseModel
+from gto.exceptions import ValidationError
 
 COMMIT = "commit"
 REF = "ref"

gto/constants.py

@@ -4,10 +4,9 @@
 from typing import Dict, Optional, Type, Union
 
 import entrypoints
+from pydantic import BaseModel
 from scmrepo.git import Git
 
-from ._pydantic import BaseModel
-
 ENRICHMENT_ENTRYPOINT = "gto.enrichment"
 
 
@@ -29,7 +28,7 @@ def get_object(self) -> BaseModel:
         raise NotImplementedError
 
     def get_dict(self):
-        return self.get_object().dict()
+        return self.get_object().model_dump()
 
     @abstractmethod
     def get_human_readable(self) -> str: