From f66cbf518968bdd1e16314a897e9ec55325dffaa Mon Sep 17 00:00:00 2001
From: vkarampudi <vkarampudi@google.com>
Date: Mon, 23 Jun 2025 09:19:58 -0700
Subject: [PATCH] Add proto wrappers

PiperOrigin-RevId: 774807442
---
 setup.py                                      |  9 ++++-
 tensorflow_model_analysis/proto/BUILD         | 34 ++++---------------
 .../proto/wrappers.proto                      | 31 +++++++++++++++++
 3 files changed, 45 insertions(+), 29 deletions(-)
 create mode 100644 tensorflow_model_analysis/proto/wrappers.proto

diff --git a/setup.py b/setup.py
index f7c5ea4a47..e464873d4c 100644
--- a/setup.py
+++ b/setup.py
@@ -102,7 +102,14 @@ def generate_proto(source, require=True):
           'or install the binary package.\n')
       sys.exit(-1)
 
-    protoc_command = [protoc, '-I../src', '-I.', '--python_out=.', source]
+    protoc_command = [
+        protoc,
+        '-I/usr/include',
+        '-I.',
+        '-I./tensorflow_model_analysis/proto',
+        '--python_out=.',
+        source,
+    ]
     if subprocess.call(protoc_command) != 0:
       sys.exit(-1)
 
diff --git a/tensorflow_model_analysis/proto/BUILD b/tensorflow_model_analysis/proto/BUILD
index bd64bd5f77..44d72eb8e9 100644
--- a/tensorflow_model_analysis/proto/BUILD
+++ b/tensorflow_model_analysis/proto/BUILD
@@ -4,40 +4,18 @@ licenses(["notice"])  # Apache 2.0
 
 package(default_visibility = ["//visibility:public"])
 
-py_proto_library(
-    name = "config_py_pb2",
-    srcs = [
-        "config.proto",
-    ],
-    default_runtime = "@com_google_protobuf//:protobuf_python",
-    protoc = "@com_google_protobuf//:protoc",
-    deps = [
-        "@com_google_protobuf//:protobuf_python",
-    ],
+filegroup(
+    name = "proto_files",
+    src = glob(["*.proto"]),
 )
 
 py_proto_library(
-    name = "metrics_for_slice_py_pb2",
-    srcs = [
-        "metrics_for_slice.proto",
-    ],
-    default_runtime = "@com_google_protobuf//:protobuf_python",
-    protoc = "@com_google_protobuf//:protoc",
-    deps = [
-        "@com_google_protobuf//:protobuf_python",
-    ],
-)
-
-py_proto_library(
-    name = "validation_result",
-    srcs = [
-        "validation_result.proto",
-    ],
+    name = "protos_py_pb2",
+    srcs = glob(["*.proto"]),
     default_runtime = "@com_google_protobuf//:protobuf_python",
     protoc = "@com_google_protobuf//:protoc",
     deps = [
-        ":config_py_pb2",
-        ":metrics_for_slice_py_pb2",
+        "__init__.py",
         "@com_google_protobuf//:protobuf_python",
     ],
 )
diff --git a/tensorflow_model_analysis/proto/wrappers.proto b/tensorflow_model_analysis/proto/wrappers.proto
new file mode 100644
index 0000000000..5a12fa13f8
--- /dev/null
+++ b/tensorflow_model_analysis/proto/wrappers.proto
@@ -0,0 +1,31 @@
+// Copyright 2018 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+syntax = "proto3";
+
+package tensorflow_model_analysis;
+
+import "google/protobuf/wrappers.proto";
+
+message MyMessage {
+  google.protobuf.DoubleValue my_double = 1;
+  google.protobuf.FloatValue my_float = 2;
+  google.protobuf.Int64Value my_int64 = 3;
+  google.protobuf.UInt64Value my_uint64 = 4;
+  google.protobuf.Int32Value my_int32 = 5;
+  google.protobuf.UInt32Value my_uint32 = 6;
+  google.protobuf.BoolValue my_bool = 7;
+  google.protobuf.StringValue my_string = 8;
+  google.protobuf.BytesValue my_bytes = 9;
+}