Coverage for tinytroupe / utils / validation.py: 33%

1import json 

2import sys 

3import unicodedata 

4 

5from pydantic import ValidationError, BaseModel 

6from tinytroupe.utils import logger 

7 

8################################################################################ 

9# Validation 

10################################################################################ 

11def check_valid_fields(obj: dict, valid_fields: list) -> None: 

12 """ 

13 Checks whether the fields in the specified dict are valid, according to the list of valid fields. If not, raises a ValueError. 

14 """ 

15 for key in obj: 

16 if key not in valid_fields: 

17 raise ValueError(f"Invalid key {key} in dictionary. Valid keys are: {valid_fields}") 

18 

19def sanitize_raw_string(value: str) -> str: 

20 """ 

21 Sanitizes the specified string by:  

22 - removing any invalid characters. 

23 - ensuring it is not longer than the maximum Python string length. 

24  

25 This is for an abundance of caution with security, to avoid any potential issues with the string. 

26 """ 

27 

28 # remove any invalid characters by making sure it is a valid UTF-8 string 

29 value = value.encode("utf-8", "ignore").decode("utf-8") 

30 

31 value = unicodedata.normalize("NFC", value) 

32 

33 

34 # ensure it is not longer than the maximum Python string length 

35 return value[:sys.maxsize] 

36 

37def sanitize_dict(value: dict) -> dict: 

38 """ 

39 Sanitizes the specified dictionary by: 

40 - removing any invalid characters. 

41 - ensuring that the dictionary is not too deeply nested. 

42 """ 

43 

44 # sanitize the string representation of the dictionary 

45 for k, v in value.items(): 

46 if isinstance(v, str): 

47 value[k] = sanitize_raw_string(v) 

48 

49 # ensure that the dictionary is not too deeply nested 

50 return value 

51 

52def to_pydantic_or_sanitized_dict(value: dict, model: BaseModel=None) -> dict: 

53 """ 

54 Converts the specified model response dictionary to a Pydantic model instance, or sanitizes it if the model is not valid. 

55 It is assumed that the dict contains the `content` key. 

56 """ 

57 

58 if model is not None and (isinstance(model, type) and issubclass(model, BaseModel)): 

59 # If a model is provided, try to validate the value against the model 

60 try: 

61 res = model.model_validate(sanitize_dict(json.loads(value['content']))) 

62 return res 

63 except ValidationError as e: 

64 logger.warning(f"Validation error: {e}") 

65 return sanitize_dict(value) 

66 else: 

67 return sanitize_dict(value) # If no model, just sanitize the dict