Merge pull request #192 from OpenFreeEnergy/pydantic-2

Use `pydantic` > 1

.github/workflows/ci-integration.yml

@@ -34,7 +34,8 @@ jobs:
     steps:
       - uses: actions/checkout@v4
 
-      - uses: mamba-org/setup-micromamba@v1
+      - name: Install environment
+        uses: mamba-org/setup-micromamba@v1
         with:
           environment-file: devtools/conda-envs/test.yml
           create-args: >-

alchemiscale/compute/api.py

@@ -109,7 +109,7 @@ def register_computeservice(
 ):
     now = datetime.utcnow()
     csreg = ComputeServiceRegistration(
-        identifier=compute_service_id, registered=now, heartbeat=now
+        identifier=ComputeServiceID(compute_service_id), registered=now, heartbeat=now
     )
 
     compute_service_id_ = n4js.register_computeservice(csreg)

alchemiscale/models.py

@@ -4,8 +4,8 @@
 
 """
 
-from typing import Optional, Union
-from pydantic import BaseModel, Field, validator, root_validator
+from typing import Optional, Union, Any
+from pydantic import BaseModel, field_validator, model_validator, ConfigDict
 from gufe.tokenization import GufeKey
 from re import fullmatch
 import unicodedata
@@ -17,6 +17,10 @@ class Scope(BaseModel):
     campaign: Optional[str] = None
     project: Optional[str] = None
 
+    model_config = ConfigDict(
+        frozen=True,
+    )
+
     def __init__(self, org=None, campaign=None, project=None):
         # we add this to allow for arg-based creation, not just keyword-based
         super().__init__(org=org, campaign=campaign, project=project)
@@ -36,9 +40,6 @@ def __eq__(self, other):
 
         return str(self) == str(other)
 
-    class Config:
-        frozen = True
-
     @staticmethod
     def _validate_component(v, component):
         """
@@ -63,20 +64,24 @@ def _validate_component(v, component):
 
         return v
 
-    @validator("org")
+    @field_validator("org")
+    @classmethod
     def valid_org(cls, v):
         return cls._validate_component(v, "org")
 
-    @validator("campaign")
+    @field_validator("campaign")
+    @classmethod
     def valid_campaign(cls, v):
         return cls._validate_component(v, "campaign")
 
-    @validator("project")
+    @field_validator("project")
+    @classmethod
     def valid_project(cls, v):
         return cls._validate_component(v, "project")
 
-    @root_validator
-    def check_scope_hierarchy(cls, values):
+    @model_validator(mode="before")
+    @classmethod
+    def check_scope_hierarchy(cls, values: Any) -> Any:
         if not _hierarchy_valid(values):
             raise InvalidScopeError(
                 f"Invalid scope hierarchy: {values}, cannot specify wildcard ('*')"
@@ -132,10 +137,10 @@ class ScopedKey(BaseModel):
     campaign: str
     project: str
 
-    class Config:
-        frozen = True
+    model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
 
-    @validator("gufe_key")
+    @field_validator("gufe_key", mode="before")
+    @classmethod
     def gufe_key_validator(cls, v):
         v = str(v)
 
@@ -157,6 +162,17 @@ def gufe_key_validator(cls, v):
         # Cast to GufeKey
         return GufeKey(v_normalized)
 
+    @model_validator(mode="before")
+    @classmethod
+    def check_scope_hierarchy(cls, values: Any) -> Any:
+        if not _hierarchy_valid(values):
+            raise InvalidScopeError(
+                f"Invalid scope hierarchy: {values}, cannot specify wildcard ('*')"
+                " in a scope component if a less specific scope component is not"
+                " given, unless all components are wildcards (*-*-*)."
+            )
+        return values
+
     def __repr__(self):  # pragma: no cover
         return f"<ScopedKey('{str(self)}')>"
 

alchemiscale/security/auth.py

@@ -72,12 +72,31 @@ def generate_secret_key():
     return secrets.token_hex(32)
 
 
-def authenticate(db, cls, identifier: str, key: str) -> CredentialedEntity:
+def authenticate(db, cls, identifier: str, key: str) -> Optional[CredentialedEntity]:
+    """Authenticate the given identity+key against the db instance.
+
+    Parameters
+    ----------
+    db
+        State store instance featuring a `get_credentialed_entity` method.
+    cls
+        The `CredentialedEntity` subclass the identity corresponds to.
+    identity
+        String identifier for the the identity.
+    key
+        Secret key string for the identity.
+
+    Returns
+    -------
+    If successfully authenticated, returns the `CredentialedEntity` subclass instance.
+    If not, returns `None`.
+
+    """
     entity: CredentialedEntity = db.get_credentialed_entity(identifier, cls)
     if entity is None:
-        return False
+        return None
     if not pwd_context.verify(key, entity.hashed_key):
-        return False
+        return None
     return entity
 
 

alchemiscale/security/models.py

@@ -7,7 +7,7 @@
 from datetime import datetime, timedelta
 from typing import List, Union, Optional
 
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, field_validator
 
 from ..models import Scope
 
@@ -32,20 +32,23 @@ class ScopedIdentity(BaseModel):
     disabled: bool = False
     scopes: List[str] = []
 
-    @validator("scopes", pre=True, each_item=True)
-    def cast_scopes_to_str(cls, scope):
+    @field_validator("scopes", mode="before")
+    @classmethod
+    def cast_scopes_to_str(cls, scopes):
         """Ensure that each scope object is correctly cast to its str representation"""
-        if isinstance(scope, Scope):
-            scope = str(scope)
-        elif isinstance(scope, str):
-            try:
-                Scope.from_str(scope)
-            except:
+        scopes_ = []
+        for scope in scopes:
+            if isinstance(scope, Scope):
+                scopes_.append(str(scope))
+            elif isinstance(scope, str):
+                try:
+                    scopes_.append(str(Scope.from_str(scope)))
+                except:
+                    raise ValueError(f"Invalid scope `{scope}` set for `{cls}`")
+            else:
                 raise ValueError(f"Invalid scope `{scope}` set for `{cls}`")
-        else:
-            raise ValueError(f"Invalid scope `{scope}` set for `{cls}`")
 
-        return scope
+        return scopes_
 
 
 class UserIdentity(ScopedIdentity):

alchemiscale/settings.py

@@ -7,12 +7,13 @@
 from functools import lru_cache
 from typing import Optional
 
-from pydantic import BaseSettings
+from pydantic_settings import BaseSettings, SettingsConfigDict
 
 
 class FrozenSettings(BaseSettings):
-    class Config:
-        frozen = True
+    model_config = SettingsConfigDict(
+        frozen=True,
+    )
 
 
 class Neo4jStoreSettings(FrozenSettings):

alchemiscale/storage/models.py

@@ -13,7 +13,7 @@
 import hashlib
 
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 from gufe.tokenization import GufeTokenizable, GufeKey
 
 from ..models import ScopedKey, Scope
@@ -29,6 +29,8 @@ class ComputeServiceRegistration(BaseModel):
     registered: datetime
     heartbeat: datetime
 
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     def __repr__(self):  # pragma: no cover
         return f"<ComputeServiceRegistration('{str(self)}')>"
 
@@ -59,6 +61,8 @@ class TaskProvenance(BaseModel):
     datetime_start: datetime
     datetime_end: datetime
 
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     # this should include versions of various libraries
 
 

devtools/conda-envs/alchemiscale-client.yml

@@ -13,7 +13,8 @@ dependencies:
   - requests
   - click
   - httpx
-  - pydantic<2.0
+  - pydantic >2
+  - pydantic-settings
   - async-lru
 
   ## user client

devtools/conda-envs/alchemiscale-compute.yml

@@ -13,7 +13,8 @@ dependencies:
   - requests
   - click
   - httpx
-  - pydantic<2.0
+  - pydantic >2
+  - pydantic-settings
   - async-lru
 
   # openmm protocols

devtools/conda-envs/alchemiscale-server.yml

@@ -13,7 +13,8 @@ dependencies:
 
   - requests
   - click
-  - pydantic<2.0
+  - pydantic >2
+  - pydantic-settings
   - async-lru
 
   ## state store

devtools/conda-envs/test.yml

@@ -9,7 +9,8 @@ dependencies:
   # alchemiscale dependencies
   - gufe>=1.1.0
   - openfe>=1.2.0
-  - pydantic<2.0
+  - pydantic >2
+  - pydantic-settings
   - async-lru
 
   ## state store

docs/conf.py

@@ -45,6 +45,7 @@
     "numpy",
     "py2neo",
     "pydantic",
+    "pydantic_settings",
     "starlette",
     "yaml",
 ]