Feedback

Author: Steve Ayers

protovalidate/internal/rules.py

@@ -23,9 +23,6 @@
 from buf.validate import validate_pb2  # type: ignore
 from protovalidate.internal.cel_field_presence import InterpretedRunner, in_has
 
-# Convenience to stringify the type names for error messages
-FIELD_TYPE_NAMES = {v: k for k, v in vars(descriptor.FieldDescriptor).items() if k.startswith("TYPE_")}
-
 
 class CompilationError(Exception):
     pass
@@ -61,59 +58,54 @@ def unwrap(msg: message.Message) -> celtypes.Value:
 }
 
 
-class MessageType(celtypes.MapType):
-    msg: message.Message
-    desc: descriptor.Descriptor
-
-    def __init__(self, msg: message.Message):
-        super().__init__()
-        self.msg = msg
-        self.desc = msg.DESCRIPTOR
-        field: descriptor.FieldDescriptor
-        for field in self.desc.fields:
-            if field.containing_oneof is not None and not self.msg.HasField(field.name):
-                continue
-            self[field.name] = field_to_cel(self.msg, field)
-
-    def __getitem__(self, name):
-        field = self.desc.fields_by_name[name]
-        if field.has_presence and not self.msg.HasField(name):
-            if in_has():
-                raise KeyError()
-            else:
-                return _zero_value(field)
-        return super().__getitem__(name)
-
-
 def _msg_to_cel(msg: message.Message) -> celtypes.Value:
     ctor = _MSG_TYPE_URL_TO_CTOR.get(msg.DESCRIPTOR.full_name)
     if ctor is not None:
         return ctor(msg)
     return MessageType(msg)
 
 
-_TYPE_TO_CTOR: dict[str, typing.Callable[..., celtypes.Value]] = {
-    descriptor.FieldDescriptor.TYPE_MESSAGE: _msg_to_cel,
-    descriptor.FieldDescriptor.TYPE_GROUP: _msg_to_cel,
-    descriptor.FieldDescriptor.TYPE_ENUM: celtypes.IntType,
-    descriptor.FieldDescriptor.TYPE_BOOL: celtypes.BoolType,
-    descriptor.FieldDescriptor.TYPE_BYTES: celtypes.BytesType,
-    descriptor.FieldDescriptor.TYPE_STRING: celtypes.StringType,
-    descriptor.FieldDescriptor.TYPE_FLOAT: celtypes.DoubleType,
-    descriptor.FieldDescriptor.TYPE_DOUBLE: celtypes.DoubleType,
-    descriptor.FieldDescriptor.TYPE_INT32: celtypes.IntType,
-    descriptor.FieldDescriptor.TYPE_INT64: celtypes.IntType,
-    descriptor.FieldDescriptor.TYPE_UINT32: celtypes.UintType,
-    descriptor.FieldDescriptor.TYPE_UINT64: celtypes.UintType,
-    descriptor.FieldDescriptor.TYPE_SINT32: celtypes.IntType,
-    descriptor.FieldDescriptor.TYPE_SINT64: celtypes.IntType,
-    descriptor.FieldDescriptor.TYPE_FIXED32: celtypes.UintType,
-    descriptor.FieldDescriptor.TYPE_FIXED64: celtypes.UintType,
-    descriptor.FieldDescriptor.TYPE_SFIXED32: celtypes.IntType,
-    descriptor.FieldDescriptor.TYPE_SFIXED64: celtypes.IntType,
+class FieldDescMetadata(typing.TypedDict):
+    name: str
+    ctor: typing.Callable[..., celtypes.Value]
+
+
+_FIELD_DESC_METADATA_MAP: dict[typing.Any, FieldDescMetadata] = {
+    descriptor.FieldDescriptor.TYPE_MESSAGE: {"name": "message", "ctor": _msg_to_cel},
+    descriptor.FieldDescriptor.TYPE_GROUP: {"name": "group", "ctor": _msg_to_cel},
+    descriptor.FieldDescriptor.TYPE_ENUM: {"name": "enum", "ctor": celtypes.IntType},
+    descriptor.FieldDescriptor.TYPE_BOOL: {"name": "bool", "ctor": celtypes.BoolType},
+    descriptor.FieldDescriptor.TYPE_BYTES: {"name": "bytes", "ctor": celtypes.BytesType},
+    descriptor.FieldDescriptor.TYPE_STRING: {"name": "string", "ctor": celtypes.StringType},
+    descriptor.FieldDescriptor.TYPE_FLOAT: {"name": "float", "ctor": celtypes.DoubleType},
+    descriptor.FieldDescriptor.TYPE_DOUBLE: {"name": "double", "ctor": celtypes.DoubleType},
+    descriptor.FieldDescriptor.TYPE_INT32: {"name": "int32", "ctor": celtypes.IntType},
+    descriptor.FieldDescriptor.TYPE_INT64: {"name": "int64", "ctor": celtypes.IntType},
+    descriptor.FieldDescriptor.TYPE_SINT32: {"name": "sint32", "ctor": celtypes.IntType},
+    descriptor.FieldDescriptor.TYPE_SINT64: {"name": "sint64", "ctor": celtypes.IntType},
+    descriptor.FieldDescriptor.TYPE_SFIXED32: {"name": "sfixed32", "ctor": celtypes.IntType},
+    descriptor.FieldDescriptor.TYPE_SFIXED64: {"name": "sfixed64", "ctor": celtypes.IntType},
+    descriptor.FieldDescriptor.TYPE_UINT32: {"name": "uint32", "ctor": celtypes.UintType},
+    descriptor.FieldDescriptor.TYPE_UINT64: {"name": "uint64", "ctor": celtypes.UintType},
+    descriptor.FieldDescriptor.TYPE_FIXED32: {"name": "fixed32", "ctor": celtypes.UintType},
+    descriptor.FieldDescriptor.TYPE_FIXED64: {"name": "fixed64", "ctor": celtypes.UintType},
 }
 
 
+def _get_type_name(fd: typing.Any) -> str:
+    md = _FIELD_DESC_METADATA_MAP.get(fd)
+    if md is None:
+        return "unknown"
+    return md["name"]
+
+
+def _get_type_ctor(fd: typing.Any) -> typing.Optional[typing.Callable[..., celtypes.Value]]:
+    md = _FIELD_DESC_METADATA_MAP.get(fd)
+    if md is None:
+        return None
+    return md["ctor"]
+
+
 def _proto_message_has_field(msg: message.Message, field: descriptor.FieldDescriptor) -> typing.Any:
     if field.is_extension:
         return msg.HasExtension(field)  # type: ignore
@@ -129,7 +121,7 @@ def _proto_message_get_field(msg: message.Message, field: descriptor.FieldDescri
 
 
 def _scalar_field_value_to_cel(val: typing.Any, field: descriptor.FieldDescriptor) -> celtypes.Value:
-    ctor = _TYPE_TO_CTOR.get(field.type)
+    ctor = _get_type_ctor(field.type)
     if ctor is None:
         msg = "unknown field type"
         raise CompilationError(msg)
@@ -234,6 +226,30 @@ def _set_path_element_map_key(
         raise CompilationError(msg)
 
 
+class MessageType(celtypes.MapType):
+    msg: message.Message
+    desc: descriptor.Descriptor
+
+    def __init__(self, msg: message.Message):
+        super().__init__()
+        self.msg = msg
+        self.desc = msg.DESCRIPTOR
+        field: descriptor.FieldDescriptor
+        for field in self.desc.fields:
+            if field.containing_oneof is not None and not self.msg.HasField(field.name):
+                continue
+            self[field.name] = field_to_cel(self.msg, field)
+
+    def __getitem__(self, name):
+        field = self.desc.fields_by_name[name]
+        if field.has_presence and not self.msg.HasField(name):
+            if in_has():
+                raise KeyError()
+            else:
+                return _zero_value(field)
+        return super().__getitem__(name)
+
+
 class Violation:
     """A singular rule violation."""
 
@@ -400,14 +416,14 @@ def check_field_type(field: descriptor.FieldDescriptor, expected: int, wrapper_n
     if field.type != expected and (
         field.type != descriptor.FieldDescriptor.TYPE_MESSAGE or field.message_type.full_name != wrapper_name
     ):
-        field_type_str = FIELD_TYPE_NAMES[field.type]
+        field_type_str = _get_type_name(field.type)
         if expected == 0:
             if wrapper_name is not None:
                 expected_type_str = wrapper_name
             else:
-                expected_type_str = FIELD_TYPE_NAMES[descriptor.FieldDescriptor.TYPE_MESSAGE]
+                expected_type_str = _get_type_name(descriptor.FieldDescriptor.TYPE_MESSAGE)
         else:
-            expected_type_str = FIELD_TYPE_NAMES[expected]
+            expected_type_str = _get_type_name(expected)
         msg = f"field {field.name} has type {field_type_str} but expected {expected_type_str}"
         raise CompilationError(msg)