Config: Switch from jsonschema to pydantic

aiida/cmdline/commands/cmd_config.py

@@ -119,7 +119,8 @@ def verdi_config_set(ctx, option, value, globally, append, remove):
 
     List values are split by whitespace, e.g. "a b" becomes ["a", "b"].
     """
-    from aiida.manage.configuration import Config, ConfigValidationError, Profile
+    from aiida.common.exceptions import ConfigurationError
+    from aiida.manage.configuration import Config, Profile
 
     if append and remove:
         echo.echo_critical('Cannot flag both append and remove')
@@ -137,7 +138,7 @@ def verdi_config_set(ctx, option, value, globally, append, remove):
     if append or remove:
         try:
             current = config.get_option(option.name, scope=scope)
-        except ConfigValidationError as error:
+        except ConfigurationError as error:
             echo.echo_critical(str(error))
         if not isinstance(current, list):
             echo.echo_critical(f'cannot append/remove to value: {current}')
@@ -149,7 +150,7 @@ def verdi_config_set(ctx, option, value, globally, append, remove):
     # Set the specified option
     try:
         value = config.set_option(option.name, value, scope=scope)
-    except ConfigValidationError as error:
+    except ConfigurationError as error:
         echo.echo_critical(str(error))
 
     config.store()

aiida/manage/__init__.py

@@ -34,7 +34,6 @@
     'BROKER_DEFAULTS',
     'CURRENT_CONFIG_VERSION',
     'Config',
-    'ConfigValidationError',
     'MIGRATIONS',
     'ManagementApiConnectionError',
     'OLDEST_COMPATIBLE_CONFIG_VERSION',
@@ -43,7 +42,6 @@
     'RabbitmqManagementClient',
     'check_and_migrate_config',
     'config_needs_migrating',
-    'config_schema',
     'disable_caching',
     'downgrade_config',
     'enable_caching',

aiida/manage/configuration/__init__.py

@@ -22,14 +22,12 @@
 __all__ = (
     'CURRENT_CONFIG_VERSION',
     'Config',
-    'ConfigValidationError',
     'MIGRATIONS',
     'OLDEST_COMPATIBLE_CONFIG_VERSION',
     'Option',
     'Profile',
     'check_and_migrate_config',
     'config_needs_migrating',
-    'config_schema',
     'downgrade_config',
     'get_current_version',
     'get_option',

aiida/manage/configuration/config.py

@@ -7,67 +7,50 @@
 # For further information on the license, see the LICENSE.txt file        #
 # For further information please visit http://www.aiida.net               #
 ###########################################################################
-"""Module that defines the configuration file of an AiiDA instance and functions to create and load it."""
+"""Module that defines the configuration file of an AiiDA instance and functions to create and load it.
+
+Despite the import of the annotations backport below which enables postponed type annotation evaluation as implemented
+with PEP 563 (https://peps.python.org/pep-0563/), this is not compatible with ``pydantic`` for Python 3.9 and older (
+See https://github.com/pydantic/pydantic/issues/2678 for details).
+"""
 from __future__ import annotations
 
 import codecs
-from functools import cache
 import json
 import os
-from typing import Any, Dict, List, Optional, Sequence, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 import uuid
 
-from pydantic import BaseModel, Field, ValidationError, validator  # pylint: disable=no-name-in-module
+from pydantic import (  # pylint: disable=no-name-in-module
+    BaseModel,
+    ConfigDict,
+    Field,
+    ValidationError,
+    field_serializer,
+    field_validator,
+)
 
 from aiida.common.exceptions import ConfigurationError
 from aiida.common.log import LogLevels
 
-from . import schema as schema_module
 from .options import Option, get_option, get_option_names, parse_option
 from .profile import Profile
 
-__all__ = ('Config', 'config_schema', 'ConfigValidationError')
-
-SCHEMA_FILE = 'config-v9.schema.json'
-
-
-@cache
-def config_schema() -> Dict[str, Any]:
-    """Return the configuration schema."""
-    from importlib.resources import files
-
-    return json.loads(files(schema_module).joinpath(SCHEMA_FILE).read_text(encoding='utf8'))
+__all__ = ('Config',)
 
 
-class ConfigValidationError(ConfigurationError):
-    """Configuration error raised when the file contents fails validation."""
-
-    def __init__(
-        self, message: str, keypath: Sequence[Any] = (), schema: Optional[dict] = None, filepath: Optional[str] = None
-    ):
-        super().__init__(message)
-        self._message = message
-        self._keypath = keypath
-        self._filepath = filepath
-        self._schema = schema
-
-    def __str__(self) -> str:
-        prefix = f'{self._filepath}:' if self._filepath else ''
-        path = '/' + '/'.join(str(k) for k in self._keypath) + ': ' if self._keypath else ''
-        schema = f'\n  schema:\n  {self._schema}' if self._schema else ''
-        return f'Validation Error: {prefix}{path}{self._message}{schema}'
-
-
-class ConfigVersionSchema(BaseModel):
+class ConfigVersionSchema(BaseModel, defer_build=True):
     """Schema for the version configuration of an AiiDA instance."""
 
     CURRENT: int
     OLDEST_COMPATIBLE: int
 
 
-class ProfileOptionsSchema(BaseModel):
+class ProfileOptionsSchema(BaseModel, defer_build=True):
     """Schema for the options of an AiiDA profile."""
 
+    model_config = ConfigDict(use_enum_values=True)
+
     runner__poll__interval: int = Field(60, description='Polling interval in seconds to be used by process runners.')
     daemon__default_workers: int = Field(
         1, description='Default number of workers to be launched by `verdi daemon start`.'
@@ -129,39 +112,37 @@ class ProfileOptionsSchema(BaseModel):
         5, description='Maximum number of transport task attempts before a Process is Paused.'
     )
     rmq__task_timeout: int = Field(10, description='Timeout in seconds for communications with RabbitMQ.')
-    storage__sandbox: Optional[str] = Field(description='Absolute path to the directory to store sandbox folders.')
+    storage__sandbox: Optional[str] = Field(
+        None, description='Absolute path to the directory to store sandbox folders.'
+    )
     caching__default_enabled: bool = Field(False, description='Enable calculation caching by default.')
     caching__enabled_for: List[str] = Field([], description='Calculation entry points to enable caching on.')
     caching__disabled_for: List[str] = Field([], description='Calculation entry points to disable caching on.')
 
-    class Config:
-        use_enum_values = True
-
-    @validator('caching__enabled_for', 'caching__disabled_for')
+    @field_validator('caching__enabled_for', 'caching__disabled_for')
     @classmethod
     def validate_caching_identifier_pattern(cls, value: List[str]) -> List[str]:
         """Validate the caching identifier patterns."""
         from aiida.manage.caching import _validate_identifier_pattern
         for identifier in value:
-            try:
-                _validate_identifier_pattern(identifier=identifier)
-            except ValueError as exception:
-                raise ValidationError(str(exception)) from exception
+            _validate_identifier_pattern(identifier=identifier)
 
         return value
 
 
 class GlobalOptionsSchema(ProfileOptionsSchema):
     """Schema for the global options of an AiiDA instance."""
-    autofill__user__email: Optional[str] = Field(description='Default user email to use when creating new profiles.')
+    autofill__user__email: Optional[str] = Field(
+        None, description='Default user email to use when creating new profiles.'
+    )
     autofill__user__first_name: Optional[str] = Field(
-        description='Default user first name to use when creating new profiles.'
+        None, description='Default user first name to use when creating new profiles.'
     )
     autofill__user__last_name: Optional[str] = Field(
-        description='Default user last name to use when creating new profiles.'
+        None, description='Default user last name to use when creating new profiles.'
     )
     autofill__user__institution: Optional[str] = Field(
-        description='Default user institution to use when creating new profiles.'
+        None, description='Default user institution to use when creating new profiles.'
     )
     rest_api__profile_switching: bool = Field(
         False, description='Toggle whether the profile can be specified in requests submitted to the REST API.'
@@ -172,14 +153,14 @@ class GlobalOptionsSchema(ProfileOptionsSchema):
     )
 
 
-class ProfileStorageConfig(BaseModel):
+class ProfileStorageConfig(BaseModel, defer_build=True):
     """Schema for the storage backend configuration of an AiiDA profile."""
 
     backend: str
     config: Dict[str, Any]
 
 
-class ProcessControlConfig(BaseModel):
+class ProcessControlConfig(BaseModel, defer_build=True):
     """Schema for the process control configuration of an AiiDA profile."""
 
     broker_protocol: str = Field('amqp', description='Protocol for connecting to the message broker.')
@@ -191,29 +172,28 @@ class ProcessControlConfig(BaseModel):
     broker_parameters: dict[str, Any] = Field('guest', description='Arguments to be encoded as query parameters.')
 
 
-class ProfileSchema(BaseModel):
+class ProfileSchema(BaseModel, defer_build=True):
     """Schema for the configuration of an AiiDA profile."""
 
     uuid: str = Field(description='', default_factory=uuid.uuid4)
     storage: ProfileStorageConfig
     process_control: ProcessControlConfig
     default_user_email: Optional[str] = None
     test_profile: bool = False
-    options: Optional[ProfileOptionsSchema]
+    options: Optional[ProfileOptionsSchema] = None
 
-    class Config:
-        json_encoders = {
-            uuid.UUID: lambda u: str(u),  # pylint: disable=unnecessary-lambda
-        }
+    @field_serializer('uuid')
+    def serialize_dt(self, value: uuid.UUID, _info):
+        return str(value)
 
 
-class ConfigSchema(BaseModel):
+class ConfigSchema(BaseModel, defer_build=True):
     """Schema for the configuration of an AiiDA instance."""
 
-    CONFIG_VERSION: Optional[ConfigVersionSchema]
-    profiles: Optional[dict[str, ProfileSchema]]
-    options: Optional[GlobalOptionsSchema]
-    default_profile: Optional[str]
+    CONFIG_VERSION: Optional[ConfigVersionSchema] = None
+    profiles: Optional[dict[str, ProfileSchema]] = None
+    options: Optional[GlobalOptionsSchema] = None
+    default_profile: Optional[str] = None
 
 
 class Config:  # pylint: disable=too-many-public-methods

aiida/manage/configuration/options.py

@@ -32,7 +32,7 @@ def name(self) -> str:
 
     @property
     def valid_type(self) -> Any:
-        return self._field.type_
+        return self._field.annotation
 
     @property
     def schema(self) -> Dict[str, Any]:
@@ -44,45 +44,51 @@ def default(self) -> Any:
 
     @property
     def description(self) -> str:
-        return self._field.field_info.description
+        return self._field.description
 
     @property
     def global_only(self) -> bool:
         from .config import ProfileOptionsSchema
-        return self._name in ProfileOptionsSchema.__fields__
+        return self._name.replace('.', '__') not in ProfileOptionsSchema.model_fields
 
     def validate(self, value: Any) -> Any:
         """Validate a value
 
         :param value: The input value
-        :param cast: Attempt to cast the value to the required type
-
         :return: The output value
-        :raise: ConfigValidationError
-
+        :raise: ConfigurationError
         """
-        value, validation_error = self._field.validate(value, {}, loc=None)
+        from pydantic import ValidationError
+
+        from .config import GlobalOptionsSchema
+
+        attribute = self.name.replace('.', '__')
 
-        if validation_error:
-            raise ConfigurationError(validation_error)
+        try:
+            result = GlobalOptionsSchema.__pydantic_validator__.validate_assignment(
+                GlobalOptionsSchema.model_construct(), attribute, value
+            )
+        except ValidationError as exception:
+            raise ConfigurationError(str(exception)) from exception
 
-        return value
+        # Return the value from the constructed model as this will have casted the value to the right type
+        return getattr(result, attribute)
 
 
 def get_option_names() -> List[str]:
     """Return a list of available option names."""
     from .config import GlobalOptionsSchema
-    return [key.replace('__', '.') for key in GlobalOptionsSchema.__fields__]
+    return [key.replace('__', '.') for key in GlobalOptionsSchema.model_fields]
 
 
 def get_option(name: str) -> Option:
     """Return option."""
     from .config import GlobalOptionsSchema
-    options = GlobalOptionsSchema.__fields__
+    options = GlobalOptionsSchema.model_fields
     option_name = name.replace('.', '__')
     if option_name not in options:
         raise ConfigurationError(f'the option {name} does not exist')
-    return Option(name, GlobalOptionsSchema.schema()['properties'][option_name], options[option_name])
+    return Option(name, GlobalOptionsSchema.model_json_schema()['properties'][option_name], options[option_name])
 
 
 def parse_option(option_name: str, option_value: Any) -> Tuple[Option, Any]:

docs/source/nitpick-exceptions

@@ -149,6 +149,15 @@ py:class ndarray
 py:class paramiko.proxy.ProxyCommand
 
 py:class pydantic.main.BaseModel
+py:class ModelPrivateAttr
+py:class CoreSchema
+py:class _decorators.DecoratorInfos
+py:class _generics.PydanticGenericMetadata
+py:class SchemaSerializer
+py:class SchemaValidator
+py:class Signature
+py:class ConfigDict
+py:class FieldInfo
 
 # These can be removed once they are properly included in the `__all__` in `plumpy`
 py:class plumpy.ports.PortNamespace

environment.yml

@@ -26,7 +26,7 @@ dependencies:
 - pgsu~=0.2.1
 - psutil~=5.6
 - psycopg2-binary~=2.8
-- pydantic~=1.10
+- pydantic~=2.4
 - pytz~=2021.1
 - pyyaml~=6.0
 - requests~=2.0

pyproject.toml

@@ -44,7 +44,7 @@ dependencies = [
     "pgsu~=0.2.1",
     "psutil~=5.6",
     "psycopg2-binary~=2.8",
-    "pydantic~=1.10",
+    "pydantic~=2.4",
     "pytz~=2021.1",
     "pyyaml~=6.0",
     "requests~=2.0",

requirements/requirements-py-3.10.txt

@@ -44,7 +44,6 @@ deprecation==2.1.0
 disk-objectstore==1.0.0
 docstring-parser==0.15
 docutils==0.16
-emmet-core==0.57.1
 exceptiongroup==1.1.1
 executing==1.2.0
 fastjsonschema==2.17.1
@@ -89,8 +88,7 @@ matplotlib-inline==0.1.6
 mdit-py-plugins==0.3.5
 mdurl==0.1.2
 mistune==3.0.1
-monty==2023.5.8
-mp-api==0.33.3
+monty==2023.9.25
 mpmath==1.3.0
 msgpack==1.0.5
 multidict==6.0.4
@@ -133,10 +131,10 @@ py-cpuinfo==9.0.0
 pybtex==0.24.0
 pycifrw==4.4.5
 pycparser==2.21
-pydantic==1.10.9
+pydantic==2.4.0
 pydata-sphinx-theme==0.13.3
 pygments==2.15.1
-pymatgen==2023.5.31
+pymatgen==2023.9.25
 pympler==0.9
 pymysql==0.9.3
 pynacl==1.5.0