Coverage for mcp_chatkit_widget / pydantic_conversion.py: 100%

1"""Helpers for translating JSON Schema definitions into Pydantic models.""" 

2 

3from typing import Any (empty)

4from pydantic import BaseModel, ConfigDict, create_model (empty)

5 

6 

7def _to_title_case(snake_or_lower: str) -> str: (empty)

8 """Convert a string to TitleCase. 

9 

10 Args: 

11 snake_or_lower: String in snake_case or lowercase 

12 (e.g., "flight_tracker" or "date") 

13 

14 Returns: 

15 String in TitleCase (e.g., "FlightTracker" or "Date") 

16 """ 

17 if "_" in snake_or_lower: (empty)

18 components = snake_or_lower.split("_") (empty)

19 return "".join(component.title() for component in components) (empty)

20 return snake_or_lower.capitalize() (empty)

21 

22 

23def _get_type_map() -> dict[str, type]: (empty)

24 """Return mapping from JSON Schema types to Python types.""" 

25 return { (empty)

26 "string": str, 

27 "integer": int, 

28 "number": float, 

29 "boolean": bool, 

30 "array": list, 

31 "object": dict, 

32 } 

33 

34 

35def _resolve_array_type( (empty)

36 field_schema: dict[str, Any], model_name: str, field_name: str 

37) -> Any: 

38 """Resolve Python type for array field schema.""" 

39 items_schema = field_schema.get("items") (empty)

40 if not isinstance(items_schema, dict): (empty)

41 return list[Any] (empty)

42 

43 item_type = items_schema.get("type") (empty)

44 if item_type == "object": (empty)

45 item_model_name = f"{model_name}{_to_title_case(field_name)}Item" (empty)

46 item_model = json_schema_to_pydantic(items_schema, item_model_name) (empty)

47 return list[item_model] # type: ignore[valid-type] (empty)

48 if item_type == "array": (empty)

49 return list[Any] (empty)

50 

51 type_map = _get_type_map() (empty)

52 mapped_type = type_map.get(item_type or "", Any) (empty)

53 return list[mapped_type] # type: ignore[valid-type] (empty)

54 

55 

56def _resolve_field_type( (empty)

57 field_schema: dict[str, Any], model_name: str, field_name: str 

58) -> Any: 

59 """Resolve Python type for a field based on its schema.""" 

60 field_type = field_schema.get("type") (empty)

61 

62 if field_type == "object": (empty)

63 nested_model_name = f"{model_name}{_to_title_case(field_name)}" (empty)

64 return json_schema_to_pydantic(field_schema, nested_model_name) (empty)

65 if field_type == "array": (empty)

66 return _resolve_array_type(field_schema, model_name, field_name) (empty)

67 

68 type_map = _get_type_map() (empty)

69 return type_map.get(field_type or "", Any) (empty)

70 

71 

72def _build_field_definitions( (empty)

73 properties: dict[str, Any], required_fields: set[str], model_name: str 

74) -> dict[str, Any]: 

75 """Build field definitions dict for Pydantic model creation.""" 

76 field_definitions: dict[str, Any] = {} (empty)

77 

78 for field_name, field_schema in properties.items(): (empty)

79 python_type = _resolve_field_type(field_schema, model_name, field_name) (empty)

80 

81 if field_name not in required_fields: (empty)

82 python_type = python_type | None (empty)

83 field_definitions[field_name] = (python_type, None) (empty)

84 else: 

85 field_definitions[field_name] = (python_type, ...) (empty)

86 

87 return field_definitions (empty)

88 

89 

90def _build_model_config( (empty)

91 schema: dict[str, Any], schema_title: str | None 

92) -> dict[str, Any]: 

93 """Build configuration kwargs for Pydantic model.""" 

94 config_kwargs: dict[str, Any] = {} (empty)

95 if schema_title: (empty)

96 config_kwargs["title"] = schema_title (empty)

97 if schema.get("additionalProperties") is False: (empty)

98 config_kwargs["extra"] = "forbid" (empty)

99 return config_kwargs (empty)

100 

101 

102def json_schema_to_pydantic( (empty)

103 schema: dict[str, Any], 

104 model_name: str = "DynamicModel", 

105 schema_title: str | None = None, 

106) -> type[BaseModel]: 

107 """Convert a JSON schema to a Pydantic model. 

108 

109 This function recursively converts JSON schema definitions into Pydantic model 

110 classes, handling nested objects and required fields. 

111 """ 

112 if schema.get("type") != "object": (empty)

113 raise ValueError("Root schema must be of type 'object'") (empty)

114 

115 properties = schema.get("properties", {}) (empty)

116 required_fields = set(schema.get("required", [])) (empty)

117 

118 field_definitions = _build_field_definitions( (empty)

119 properties, required_fields, model_name 

120 ) 

121 config_kwargs = _build_model_config(schema, schema_title) (empty)

122 

123 if config_kwargs: (empty)

124 config = ConfigDict(**config_kwargs) # type: ignore[typeddict-item] (empty)

125 return create_model(model_name, __config__=config, **field_definitions) (empty)

126 

127 return create_model(model_name, **field_definitions) (empty)

128 

129 

130__all__ = [ (empty)

131 "json_schema_to_pydantic", 

132 "_to_title_case", 

133]