Tensorflow protobuf stubs (#9873)

Co-authored-by: Mehdi Drissi <mdrissi@snapchat.com>
Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>
Co-authored-by: Jelle Zijlstra <jelle.zijlstra@gmail.com>

scripts/sync_tensorflow_protobuf_stubs.sh

@@ -0,0 +1,65 @@
+#!/bin/bash
+set -euxo pipefail
+
+# Partly based on scripts/generate_proto_stubs.sh.
+
+# Generates the protobuf stubs for the given tensorflow version using
+# mypy-protobuf. Should be run like ./sync_tensorflow_protobuf_stubs.sh
+# Generally, new minor versions are a good time to update the stubs.
+cd "$(dirname "$0")" > /dev/null
+cd ../stubs/tensorflow
+REPO_ROOT="$(realpath "$(dirname "${BASH_SOURCE[0]}")"/..)"
+
+
+# This version should be consistent with the version in tensorflow's METADATA.toml.
+TENSORFLOW_VERSION=2.11.0
+# Latest mypy-protobuf has dependency on protobuf >4, which is incompatible at runtime
+# with tensorflow. However, the stubs produced do still work with tensorflow. So after
+# installing mypy-protobuf, before running stubtest on tensorflow you should downgrade
+# protobuf<4.
+MYPY_PROTOBUF_VERSION=3.4.0
+
+pip install mypy-protobuf=="$MYPY_PROTOBUF_VERSION"
+
+mkdir repository
+pushd repository &> /dev/null
+    git clone https://github.com/tensorflow/tensorflow.git
+    pushd tensorflow &> /dev/null
+        git checkout v"$TENSORFLOW_VERSION"
+
+        # Folders here cover the more commonly used protobufs externally and
+        # their dependencies. Tensorflow has more protobufs and can be added if requested.
+        protoc --mypy_out "relax_strict_optional_primitives:$REPO_ROOT/stubs/tensorflow" \
+            tensorflow/core/protobuf/*.proto \
+            tensorflow/core/protobuf/tpu/*.proto \
+            tensorflow/core/framework/*.proto \
+            tensorflow/core/util/*.proto \
+            tensorflow/core/example/*.proto \
+            tensorflow/python/keras/protobuf/*.proto \
+            tensorflow/tsl/protobuf/*.proto \
+            tensorflow/compiler/xla/*.proto \
+            tensorflow/compiler/xla/service/*.proto
+    popd &> /dev/null
+popd &> /dev/null
+
+rm -rf repository/
+# These protos exist in a folder with protos used in python, but are not
+# included in the python wheel. They are likely only used for other
+# language builds. stubtest was used to identify them by looking for
+# ModuleNotFoundError.
+rm tensorflow/core/protobuf/coordination_service_pb2.pyi \
+   tensorflow/compiler/xla/service/hlo_execution_profile_data_pb2.pyi \
+   tensorflow/compiler/xla/service/hlo_profile_printer_data_pb2.pyi \
+   tensorflow/compiler/xla/service/test_compilation_environment_pb2.pyi \
+   tensorflow/core/protobuf/autotuning_pb2.pyi \
+   tensorflow/core/protobuf/conv_autotuning_pb2.pyi \
+   tensorflow/core/protobuf/critical_section_pb2.pyi \
+   tensorflow/core/protobuf/eager_service_pb2.pyi \
+   tensorflow/core/protobuf/master_pb2.pyi \
+   tensorflow/core/protobuf/master_service_pb2.pyi \
+   tensorflow/core/protobuf/replay_log_pb2.pyi \
+   tensorflow/core/protobuf/worker_pb2.pyi \
+   tensorflow/core/protobuf/worker_service_pb2.pyi \
+   tensorflow/core/protobuf/tpu/compile_metadata_pb2.pyi \
+   tensorflow/core/util/example_proto_fast_parsing_test_pb2.pyi \
+   tensorflow/compiler/xla/xla_pb2.pyi

stubs/tensorflow/@tests/stubtest_allowlist.txt

@@ -64,3 +64,7 @@ tensorflow.keras.layers.Layer.__new__
 # disagreements.
 tensorflow.keras.layers.*.build
 tensorflow.keras.layers.*.compute_output_shape
+
+# pb2.pyi generated by mypy-protobuf diverge with runtime in many ways. These stubs
+# are mainly tested in mypy-protobuf.
+.*_pb2.*

stubs/tensorflow/METADATA.toml

@@ -1,6 +1,6 @@
 version = "2.11.*"
 # requires a version of numpy with a `py.typed` file
-requires = ["numpy>=1.20"]
+requires = ["numpy>=1.20", "types-protobuf"]
 
 [tool.stubtest]
 ignore_missing_stub = true