Coverage for src/bioimageio/spec/_internal/validator_annotations.py: 85%

1from __future__ import annotations 

2 

3from dataclasses import dataclass 

4from typing import Any 

5 

6import annotated_types 

7from pydantic import GetCoreSchemaHandler, functional_validators 

8from pydantic_core import CoreSchema 

9from pydantic_core.core_schema import no_info_after_validator_function 

10 

11from .utils import SLOTS 

12 

13 

14# TODO: make sure we use this one everywhere and not the vanilla pydantic one 

15@dataclass(frozen=True, **SLOTS) 

16class AfterValidator(functional_validators.AfterValidator): 

17 def __str__(self): 

18 return f"AfterValidator({self.func.__name__})" 

19 

20 

21# TODO: make sure we use this one everywhere and not the vanilla pydantic one 

22@dataclass(frozen=True, **SLOTS) 

23class BeforeValidator(functional_validators.BeforeValidator): 

24 def __str__(self): 

25 return f"BeforeValidator({self.func.__name__})" 

26 

27 

28# TODO: make sure we use this one everywhere and not the vanilla pydantic one 

29@dataclass(frozen=True, **SLOTS) 

30class Predicate(annotated_types.Predicate): 

31 def __str__(self): 

32 return f"Predicate({self.func.__name__})" 

33 

34 

35@dataclass(frozen=True, **SLOTS) 

36class RestrictCharacters: 

37 alphabet: str 

38 

39 def __get_pydantic_core_schema__( 

40 self, source: type[Any], handler: GetCoreSchemaHandler 

41 ) -> CoreSchema: 

42 if not self.alphabet: 

43 raise ValueError("Alphabet may not be empty") 

44 

45 schema = handler(source) # get the CoreSchema from the type / inner constraints 

46 if schema["type"] != "str" and not ( 

47 schema["type"] == "function-after" and schema["schema"]["type"] == "str" 

48 ): 

49 raise TypeError("RestrictCharacters can only be applied to strings") 

50 

51 return no_info_after_validator_function( 

52 self.validate, 

53 schema, 

54 ) 

55 

56 def validate(self, value: str) -> str: 

57 if any(c not in self.alphabet for c in value): 

58 raise ValueError(f"{value!r} is not restricted to {self.alphabet!r}") 

59 return value