def test_gather_file_descriptors(self):
fd_names = set()
for fd in proto_utils.gather_file_descriptors(foo_pb2.Foo.DESCRIPTOR):
fd_names.add(fd.name)
self.assertEqual(
fd_names,
{'tfx/utils/testdata/bar.proto', 'tfx/utils/testdata/foo.proto'})
def encode(
self,
sub_expression_pb: placeholder_pb2.PlaceholderExpression,
component_spec: Optional[types.ComponentSpec] = None
) -> placeholder_pb2.PlaceholderExpression:
result = placeholder_pb2.PlaceholderExpression()
result.operator.proto_op.expression.CopyFrom(sub_expression_pb)
result.operator.proto_op.proto_field_path.extend(self._proto_field_path)
# Attach proto descriptor if available through component spec.
if (component_spec and sub_expression_pb.placeholder.type ==
placeholder_pb2.Placeholder.EXEC_PROPERTY):
exec_property_name = sub_expression_pb.placeholder.key
if exec_property_name not in component_spec.PARAMETERS:
raise ValueError(
f"Can't find provided placeholder key {exec_property_name} in "
"component spec's exec properties. "
f"Available exec property keys: {component_spec.PARAMETERS.keys()}."
)
execution_param = component_spec.PARAMETERS[exec_property_name]
if not issubclass(execution_param.type, message.Message):
raise ValueError(
"Can't apply placehodler proto operator on non-proto type "
f"exec property. Got {execution_param.type}.")
fd_set = result.operator.proto_op.proto_schema.file_descriptors
for fd in proto_utils.gather_file_descriptors(
execution_param.type.DESCRIPTOR):
fd.CopyToProto(fd_set.file.add())
return result
