support pydantic v2

gto/_pydantic.py

@@ -0,0 +1,30 @@
+__all__ = [
+    "BaseModel",
+    "BaseSettings",
+    "ValidationError",
+    "parse_obj_as",
+    "validator",
+    "InitSettingsSource",
+]
+
+
+try:
+    from pydantic.v1 import (
+        BaseModel,
+        BaseSettings,
+        ValidationError,
+        parse_obj_as,
+        validator,
+    )
+    from pydantic.v1.env_settings import InitSettingsSource
+except ImportError:
+    from pydantic import (  # type: ignore[no-redef,assignment]
+        BaseModel,
+        BaseSettings,
+        ValidationError,
+        parse_obj_as,
+        validator,
+    )
+    from pydantic.env_settings import (  # type: ignore[no-redef,assignment]
+        InitSettingsSource,
+    )

gto/base.py

@@ -1,7 +1,6 @@
 from datetime import datetime
 from typing import Any, Dict, FrozenSet, List, Optional, Sequence, Union
 
-from pydantic import BaseModel
 from scmrepo.git import Git
 
 from gto.config import RegistryConfig
@@ -13,6 +12,7 @@
 )
 from gto.versions import SemVer
 
+from ._pydantic import BaseModel
 from .exceptions import (
     ArtifactNotFound,
     ManyVersions,

gto/config.py

@@ -3,14 +3,14 @@
 from pathlib import Path
 from typing import Any, Dict, List, Optional
 
-from pydantic import BaseModel, BaseSettings, validator
-from pydantic.env_settings import InitSettingsSource
 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,8 +27,8 @@ def load(self) -> EnrichmentReader:
 
 class NoFileConfig(BaseSettings):  # type: ignore[valid-type]
     INDEX: str = "artifacts.yaml"
-    TYPES: Optional[List[str]]
-    STAGES: Optional[List[str]]
+    TYPES: Optional[List[str]] = None
+    STAGES: Optional[List[str]] = None
     LOG_LEVEL: str = "INFO"
     DEBUG: bool = False
     ENRICHMENTS: List[EnrichmentConfig] = []
@@ -41,11 +41,13 @@ class Config:
 
     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)
 
     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)
 

gto/constants.py

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

gto/ext.py

@@ -4,9 +4,10 @@
 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"
 
 

gto/index.py

@@ -19,7 +19,6 @@
     Union,
 )
 
-from pydantic import BaseModel, ValidationError, parse_obj_as, validator
 from ruamel.yaml import YAMLError
 from scmrepo.exceptions import SCMError
 from scmrepo.git import Git
@@ -51,6 +50,8 @@
 from gto.git_utils import RemoteRepoMixin
 from gto.ui import echo
 
+from ._pydantic import BaseModel, ValidationError, parse_obj_as, validator
+
 logger = logging.getLogger("gto")
 
 

gto/registry.py

@@ -4,7 +4,6 @@
 from typing import List, Optional, TypeVar, cast
 
 from funcy import distinct
-from pydantic import BaseModel
 from scmrepo.git import Git
 
 from gto.base import (
@@ -37,6 +36,8 @@
 from gto.ui import echo
 from gto.versions import SemVer
 
+from ._pydantic import BaseModel
+
 TBaseEvent = TypeVar("TBaseEvent", bound=BaseEvent)
 
 

gto/tag.py

@@ -4,10 +4,10 @@
 from enum import Enum
 from typing import FrozenSet, Iterable, Optional, Union
 
-from pydantic import BaseModel
 from scmrepo.exceptions import RevError
 from scmrepo.git import Git, GitTag
 
+from ._pydantic import BaseModel
 from .base import (
     Artifact,
     Assignment,

gto/utils.py

@@ -6,11 +6,12 @@
 from enum import Enum
 
 import click
-from pydantic import BaseModel
 from tabulate import tabulate
 
 from gto.config import yaml
 
+from ._pydantic import BaseModel
+
 
 def flatten(obj):
     if isinstance(obj, Iterable) and not isinstance(obj, (str, bytes)):

setup.py

@@ -6,7 +6,9 @@
     "scmrepo>=1.3.1,<2",
     "typer>=0.4.1",
     "rich",
-    "pydantic>=1.9.0,<2",
+    # pydantic.v1.parse_obj is broken in ==2.0.0:
+    # https://github.com/pydantic/pydantic/issues/6361
+    "pydantic>=1.9.0,!=2.0.0",
     "ruamel.yaml",
     "semver>=3.0.0",
     "entrypoints",

tests/utils.py

@@ -3,7 +3,8 @@
 from typing import Any, Dict, Sequence, Set, Union
 
 from funcy import omit
-from pydantic import BaseModel
+
+from gto._pydantic import BaseModel
 
 
 def show_difference(left: Dict, right: Dict):