Pydantic v2

jupyterhub_moss/models.py

@@ -4,32 +4,29 @@
 from typing import Dict, Optional
 
 from pydantic import (
+    constr,
+    field_validator,
     BaseModel,
-    ConstrainedStr,
-    Extra,
+    ConfigDict,
+    FieldValidationInfo,
     NonNegativeInt,
     PositiveInt,
-    validator,
+    RootModel,
 )
 
-# constrained types and validators
+# Validators
 
 
-class NonEmptyStr(ConstrainedStr):
-    min_length = 1
-    strip_whitespace = True
-
-
-def check_match_gpu(v: Optional[int], values: dict) -> Optional[int]:
-    if v is not None and v > 0 and values.get("gpu") == "":
+def check_match_gpu(v: Optional[int], info: FieldValidationInfo) -> Optional[int]:
+    if v is not None and v > 0 and info.data.get("gpu") == "":
         return 0  # GPU explicitly disabled
     return v
 
 
 # models
 
 
-class PartitionResources(BaseModel, allow_mutation=False, extra=Extra.allow):
+class PartitionResources(BaseModel, frozen=True, extra="allow"):
     """SLURM partition required resources information
 
     This information retrieved from SLURM is used to constraint user's selection.
@@ -44,12 +41,10 @@ class PartitionResources(BaseModel, allow_mutation=False, extra=Extra.allow):
     max_runtime: PositiveInt
 
     # validators
-    _check_match_gpu = validator("max_ngpus", allow_reuse=True)(check_match_gpu)
+    _check_match_gpu = field_validator("max_ngpus")(check_match_gpu)
 
 
-class PartitionAllResources(
-    PartitionResources, allow_mutation=False, extra=Extra.forbid
-):
+class PartitionAllResources(PartitionResources, frozen=True, extra="forbid"):
     """SLURM partition resources information
 
     Extends resource constraints information with information
@@ -62,41 +57,40 @@ class PartitionAllResources(
     ncores_idle: NonNegativeInt
 
 
-class JupyterEnvironment(BaseModel, allow_mutation=False, extra=Extra.forbid):
+class JupyterEnvironment(BaseModel, frozen=True, extra="forbid"):
     """Single Jupyter environement description"""
 
-    add_to_path = True
-    description: NonEmptyStr
-    path = ""
-    modules = ""
-    prologue = ""
+    add_to_path: bool = True
+    # See https://github.com/pydantic/pydantic/issues/156 for type: ignore
+    description: constr(strip_whitespace=True, min_length=1)  # type: ignore[valid-type]
+    path: str = ""
+    modules: str = ""
+    prologue: str = ""
 
     # validators
-    @validator("modules")
-    def check_path_or_mods(cls, v: str, values: dict) -> str:
-        if not v and not values.get("path"):
+    @field_validator("modules")
+    def check_path_or_mods(cls, v: str, info: FieldValidationInfo) -> str:
+        if not v and not info.data.get("path"):
             raise ValueError("Jupyter environment path or modules is required")
         return v
 
 
-class PartitionConfig(BaseModel, allow_mutation=False, extra=Extra.forbid):
+class PartitionConfig(BaseModel, frozen=True, extra="forbid"):
     """Information about partition description and available environments"""
 
-    architecture = ""
-    description = ""
+    architecture: str = ""
+    description: str = ""
     jupyter_environments: Dict[str, JupyterEnvironment]
-    simple = True
+    simple: bool = True
 
 
-class PartitionInfo(
-    PartitionConfig, PartitionResources, allow_mutation=False, extra=Extra.allow
-):
+class PartitionInfo(PartitionConfig, PartitionResources):
     """Complete information about a partition: config and resources"""
 
-    pass
+    model_config = ConfigDict(frozen=True, extra="allow")
 
 
-class _PartitionTraits(PartitionConfig, allow_mutation=False, extra=Extra.forbid):
+class _PartitionTraits(PartitionConfig, frozen=True, extra="forbid"):
     """Configuration of a single partition passed as ``partitions`` traits"""
 
     gpu: Optional[str] = None
@@ -105,53 +99,59 @@ class _PartitionTraits(PartitionConfig, allow_mutation=False, extra=Extra.forbid
     max_runtime: Optional[int] = None
 
     # validators
-    _check_match_gpu = validator("max_ngpus", allow_reuse=True)(check_match_gpu)
+    _check_match_gpu = field_validator("max_ngpus")(check_match_gpu)
 
-    @validator("max_ngpus")
+    @field_validator("max_ngpus")
     def check_is_positive_or_none(cls, v: Optional[int]) -> Optional[int]:
         if v is not None and v < 0:
             raise ValueError("Value must be positive")
         return v
 
-    @validator("max_nprocs", "max_runtime")
+    @field_validator("max_nprocs", "max_runtime")
     def check_is_strictly_positive_or_none(cls, v: Optional[int]) -> Optional[int]:
         if v is not None and v <= 0:
             raise ValueError("Value must be strictly positive")
         return v
 
 
-class PartitionsTrait(BaseModel, allow_mutation=False, extra=Extra.forbid):
+class PartitionsTrait(RootModel):
     """Configuration passed as ``partitions`` trait"""
 
-    __root__: Dict[str, _PartitionTraits]
+    root: Dict[str, _PartitionTraits]
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
 
-    def dict(self, *args, **kwargs):
-        return {k: v.dict(*args, **kwargs) for k, v in self.__root__.items()}
+    def model_dump(self, *args, **kwargs):
+        return {k: v.model_dump(*args, **kwargs) for k, v in self.root.items()}
 
     def items(self):
-        return self.__root__.items()
+        return self.root.items()
+
+
+_MEM_REGEXP = re.compile("^[0-9]*([0-9]+[KMGT])?$")
+"""Memory input regular expression"""
 
 
 class UserOptions(BaseModel):
     """Options passed as `Spawner.user_options`"""
 
     # Options received through the form or GET request
     partition: str
-    runtime = ""
+    runtime: str = ""
     nprocs: PositiveInt = 1
-    memory = ""
-    reservation = ""
+    memory: str = ""
+    reservation: str = ""
     ngpus: NonNegativeInt = 0
-    options = ""
-    output = "/dev/null"
-    environment_id = ""
-    environment_path = ""
-    environment_modules = ""
-    default_url = ""
-    root_dir = ""
+    options: str = ""
+    output: str = "/dev/null"
+    environment_id: str = ""
+    environment_path: str = ""
+    environment_modules: str = ""
+    default_url: str = ""
+    root_dir: str = ""
     # Extra fields
-    gres = ""
-    prologue = ""
+    gres: str = ""
+    prologue: str = ""
 
     @classmethod
     def parse_formdata(cls, formdata: dict[str, list[str]]) -> UserOptions:
@@ -169,9 +169,9 @@ def parse_formdata(cls, formdata: dict[str, list[str]]) -> UserOptions:
         fields["output"] = (
             "slurm-%j.out" if fields.get("output", "false") == "true" else "/dev/null"
         )
-        return cls.parse_obj(fields)
+        return cls.model_validate(fields)
 
-    @validator(
+    @field_validator(
         "partition",
         "runtime",
         "memory",
@@ -189,24 +189,22 @@ def has_no_newline(cls, v: str) -> str:
             raise ValueError("Must not contain newline")
         return v
 
-    @validator("default_url")
+    @field_validator("default_url")
     def is_absolute_path(cls, v: str) -> str:
         if v and not v.startswith("/"):
             raise ValueError("Must start with /")
         return v
 
-    @validator("runtime")
+    @field_validator("runtime")
     def check_timelimit(cls, v: str) -> str:
         from .utils import parse_timelimit  # avoid circular imports
 
         if v:
             parse_timelimit(v)  # Raises exception if malformed
         return v
 
-    _MEM_REGEXP = re.compile("^[0-9]*([0-9]+[KMGT])?$")
-
-    @validator("memory")
+    @field_validator("memory")
     def check_memory(cls, v: str) -> str:
-        if v and cls._MEM_REGEXP.match(v) is None:
+        if v and _MEM_REGEXP.match(v) is None:
             raise ValueError("Error in memory syntax")
         return v

jupyterhub_moss/spawner.py

@@ -74,7 +74,7 @@ class MOSlurmSpawner(SlurmSpawner):
 
     @traitlets.validate("partitions")
     def _validate_partitions(self, proposal: dict) -> dict[str, dict]:
-        return PartitionsTrait.parse_obj(proposal["value"]).dict()
+        return PartitionsTrait.model_validate(proposal["value"]).model_dump()
 
     slurm_info_cmd = traitlets.Unicode(
         # Get number of nodes/state, cores/node, cores/state, gpus, total memory for all partitions
@@ -222,14 +222,14 @@ async def _get_partitions_info(self) -> dict[str, PartitionInfo]:
         resources_info = self.slurm_info_resources(out)
         self.log.debug("Slurm partition resources: %s", resources_info)
 
-        partitions = PartitionsTrait.parse_obj(self.partitions)
+        partitions = PartitionsTrait.model_validate(self.partitions)
 
         # use data from Slurm as base and overwrite with manual configuration settings
         partitions_info = {
-            partition: PartitionInfo.parse_obj(
+            partition: PartitionInfo.model_validate(
                 {
-                    **resources_info[partition].dict(),
-                    **config_partition_info.dict(exclude_none=True),
+                    **resources_info[partition].model_dump(),
+                    **config_partition_info.model_dump(exclude_none=True),
                 }
             )
             for partition, config_partition_info in partitions.items()
@@ -251,7 +251,7 @@ async def create_options_form(spawner: MOSlurmSpawner) -> str:
         # Strip prologue from partitions_info:
         # it is not useful and can cause some parsing issues
         partitions_dict = {
-            name: info.dict(
+            name: info.model_dump(
                 exclude={
                     "jupyter_environments": {
                         env_name: {"prologue"} for env_name in info.jupyter_environments
@@ -342,7 +342,7 @@ async def options_from_form(self, formdata: dict[str, list[str]]) -> dict:
         self.__validate_options(options, partition_info)
         self.__update_options(options, partition_info)
 
-        return options.dict()
+        return options.model_dump()
 
     def __update_spawn_commands(self, cmd_path: str) -> None:
         """Add path to commands"""

setup.cfg

@@ -26,7 +26,7 @@ install_requires =
     batchspawner>=1.0
     jinja2
     jupyterhub
-    pydantic
+    pydantic>=2.0,<3
     traitlets
 
 [options.extras_require]