Coverage for apio/common/proto_util.py: 91%

1"""Utilities related to the Apio Protocol Buffers objects.""" 

2 

3import re 

4from typing import Any, Dict, TypeVar 

5from google.protobuf.json_format import ParseDict 

6from google.protobuf.message import Message 

7from google.protobuf.unknown_fields import UnknownFieldSet 

8from google.protobuf.json_format import MessageToDict 

9from apio.common.apio_console import fatal_error 

10 

11# Placeholder for a concrete protobuf message *class* (e.g. MyProto), 

12# not an instance. bound=Message restricts it to subclasses of 

13# google.protobuf.message.Message. A TypeVar (vs plain type[Message]) 

14# lets the type checker keep the specific class: pass MyProto, get MyProto. 

15MessageClass = TypeVar("MessageClass", bound=Message) 

16 

17 

18def check_is_initialized( 

19 proto_msg: Message, error_context: str, *, json_naming: bool = False 

20) -> None: 

21 """Check that a proto message is fully populated""" 

22 

23 assert isinstance(proto_msg, Message), type(proto_msg) 

24 

25 # -- Check 1: All required fields should present. 

26 if not proto_msg.IsInitialized(): 

27 # -- Report the first missing required field. 

28 # missing_field: str = proto_msg.FindInitializationErrors()[0] 

29 find = getattr(proto_msg, "FindInitializationErrors", None) 

30 assert callable(find) 

31 missing_field = str(find()[0]) 

32 if json_naming: 32 ↛ 34line 32 didn't jump to line 34 because the condition on line 32 was always true

33 missing_field = missing_field.replace("_", "-") 

34 fatal_error(error_context, f"Missing required field '{missing_field}'") 

35 

36 # -- Check 2: Should not carry unknown fields. 

37 unknown_fields = list(UnknownFieldSet(proto_msg)) 

38 if len(unknown_fields) > 0: 38 ↛ 39line 38 didn't jump to line 39 because the condition on line 38 was never true

39 fatal_error( 

40 error_context, f'Unknown fields: {", ".join(unknown_fields)}' 

41 ) 

42 

43 

44def check_is_required(proto_msg: Message, *fields_names: str) -> None: 

45 """Check that all the names are of required fields of the proto 

46 object proto_msg. Names may be nested, e.g. "field1.field2". 

47 """ 

48 assert isinstance(proto_msg, Message), type(proto_msg) 

49 

50 for name in fields_names: 

51 descriptor = proto_msg.DESCRIPTOR 

52 for segment in name.split("."): 

53 field = descriptor.fields_by_name.get(segment) 

54 if field is None: 

55 fatal_error( 

56 f"Field '{name}' is not a field of " 

57 "protocol buffer message " 

58 f"'{proto_msg.DESCRIPTOR.full_name}'" 

59 ) 

60 if not field.is_required: 

61 fatal_error( 

62 f"Field '{name}' of '{proto_msg.DESCRIPTOR.full_name}' " 

63 f"is not required" 

64 ) 

65 descriptor = field.message_type 

66 

67 

68def check_not_required(proto_msg: Message, *fields_names: str) -> None: 

69 """Check that none of the names is a not fully required path of proto_msg. 

70 Names may be nested, e.g. "field1.field2". A path is not required if 

71 any segment along it is not required. So if "b" is optional, "a.b.c" 

72 is not required. Repeated fields are not required since they can have 

73 zero members. 

74 """ 

75 assert isinstance(proto_msg, Message), type(proto_msg) 

76 

77 for name in fields_names: 

78 descriptor = proto_msg.DESCRIPTOR 

79 all_required = True 

80 for segment in name.split("."): 

81 if descriptor is None: 81 ↛ 82line 81 didn't jump to line 82 because the condition on line 81 was never true

82 fatal_error( 

83 f"Field '{name}' is not a field of " 

84 "protocol buffer message " 

85 f"'{proto_msg.DESCRIPTOR.full_name}'" 

86 ) 

87 field = descriptor.fields_by_name.get(segment) 

88 if field is None: 88 ↛ 89line 88 didn't jump to line 89 because the condition on line 88 was never true

89 fatal_error( 

90 f"Field '{name}' is not a field of " 

91 "protocol buffer message " 

92 f"'{proto_msg.DESCRIPTOR.full_name}'" 

93 ) 

94 if not field.is_required: 

95 all_required = False 

96 descriptor = field.message_type 

97 if all_required: 

98 fatal_error( 

99 f"Field '{name}' of '{proto_msg.DESCRIPTOR.full_name}' " 

100 f"is required" 

101 ) 

102 

103 

104def proto_from_json_dict( 

105 json_dict: Dict[str, Any], 

106 proto_class: type[MessageClass], 

107 error_context: str, 

108) -> MessageClass: 

109 """Create and return an object of proto message class 'proto_class' 

110 populated with values from json dict json_dict. Exit with an error code 

111 on any error. 

112 """ 

113 # pylint: disable=broad-exception-caught 

114 

115 try: 

116 proto_msg = ParseDict(json_dict, proto_class()) 

117 except Exception as e: 

118 error_msg = str(e) 

119 

120 # -- Try to improve the error message. 

121 pattern = re.compile(r'has no field named "([^"]+)" at') 

122 match = pattern.search(error_msg) 

123 if match: 123 ↛ 125line 123 didn't jump to line 125 because the condition on line 123 was always true

124 error_msg = f"Unknown field '{match.group(1)}'" 

125 fatal_error(error_context, error_msg) 

126 

127 check_is_initialized(proto_msg, error_context, json_naming=True) 

128 return proto_msg 

129 

130 

131def proto_to_json_dict(proto_msg: Message) -> Dict[str, Any]: 

132 """Given a proto object, convert it to a json dict.""" 

133 assert isinstance(proto_msg, Message), type(proto_msg) 

134 return MessageToDict(proto_msg)