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

1from dataclasses import dataclass 

2from typing import Any, Type 

3 

4import annotated_types 

5from pydantic import GetCoreSchemaHandler, functional_validators 

6from pydantic_core import CoreSchema 

7from pydantic_core.core_schema import no_info_after_validator_function 

8 

9from .utils import SLOTS 

10 

11 

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

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

14class AfterValidator(functional_validators.AfterValidator): 

15 def __str__(self): 

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

17 

18 

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

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

21class BeforeValidator(functional_validators.BeforeValidator): 

22 def __str__(self): 

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

24 

25 

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

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

28class Predicate(annotated_types.Predicate): 

29 def __str__(self): 

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

31 

32 

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

34class RestrictCharacters: 

35 alphabet: str 

36 

37 def __get_pydantic_core_schema__( 

38 self, source: Type[Any], handler: GetCoreSchemaHandler 

39 ) -> CoreSchema: 

40 if not self.alphabet: 

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

42 

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

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

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

46 ): 

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

48 

49 return no_info_after_validator_function( 

50 self.validate, 

51 schema, 

52 ) 

53 

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

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

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

57 return value