mypy-protobuf v3.5: Fix and run protobuf generation scripts

scripts/generate_proto_stubs.sh

@@ -13,7 +13,7 @@ set -ex -o pipefail
 # Update these two variables when rerunning script
 PROTOBUF_VERSION=21.8
 PYTHON_PROTOBUF_VERSION=4.21.8
-MYPY_PROTOBUF_VERSION=v3.5.0
+MYPY_PROTOBUF_VERSION=3.5.0
 
 if uname -a | grep Darwin; then
     # brew install coreutils wget
@@ -48,7 +48,7 @@ source "$VENV/bin/activate"
 pip install -r "$REPO_ROOT/requirements-tests.txt"  # for black and isort
 
 # Install mypy-protobuf
-pip install "git+https://github.com/dropbox/mypy-protobuf@$MYPY_PROTOBUF_VERSION"
+pip install mypy-protobuf=="$MYPY_PROTOBUF_VERSION"
 
 # Remove existing pyi
 find "$REPO_ROOT/stubs/protobuf/" -name '*_pb2.pyi' -delete
@@ -76,5 +76,6 @@ protoc_install/bin/protoc --proto_path="$PYTHON_PROTOBUF_DIR/src" --mypy_out="re
 isort "$REPO_ROOT/stubs/protobuf"
 black "$REPO_ROOT/stubs/protobuf"
 
-sed --in-place="" "s/mypy-protobuf [^\"]*/mypy-protobuf ${MYPY_PROTOBUF_VERSION}/" "$REPO_ROOT/stubs/protobuf/METADATA.toml"
-sed --in-place="" "s/version = .*$/version = \"$(echo ${PYTHON_PROTOBUF_VERSION} | cut -d. -f1-2)\.\*\"/" "$REPO_ROOT/stubs/protobuf/METADATA.toml"
+sed --in-place="" \
+  "s/extra_description = .*$/extra_description = \"Generated using [mypy-protobuf==$MYPY_PROTOBUF_VERSION](https:\/\/github.com\/nipunn1313\/mypy-protobuf\/tree\/v$MYPY_PROTOBUF_VERSION) on protobuf==$PYTHON_PROTOBUF_VERSION\"/" \
+  "$REPO_ROOT/stubs/protobuf/METADATA.toml"

scripts/sync_tensorflow_protobuf_stubs.sh

@@ -3,63 +3,70 @@ 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
+# Generates the protobuf stubs for the given tensorflow version using mypy-protobuf.
 # 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
+TENSORFLOW_VERSION=2.12.1
 # 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
+MYPY_PROTOBUF_VERSION=3.5.0
 
 pip install mypy-protobuf=="$MYPY_PROTOBUF_VERSION"
 
-mkdir repository
+cd "$(dirname "$0")" > /dev/null
+cd ../stubs/tensorflow
+mkdir -p repository
 pushd repository &> /dev/null
-    git clone https://github.com/tensorflow/tensorflow.git
+    # If the script fails halfway, it's nice to be able to re-run it immediately
+    if [ ! -d "tensorflow" ] ; then
+        git clone --depth 1 --branch v"$TENSORFLOW_VERSION" https://github.com/tensorflow/tensorflow.git
+    fi
     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/compiler/xla/*.proto \
+            tensorflow/compiler/xla/service/*.proto \
+            tensorflow/core/example/*.proto \
+            tensorflow/core/framework/*.proto \
             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
+            tensorflow/tsl/protobuf/*.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 \
+rm 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/compiler/xla/xla_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/tpu/compile_metadata_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
+   tensorflow/core/util/example_proto_fast_parsing_test_pb2.pyi
+
+isort "$REPO_ROOT/stubs/tensorflow/tensorflow"
+black "$REPO_ROOT/stubs/tensorflow/tensorflow"
+
+sed --in-place="" \
+    "s/extra_description = .*$/extra_description = \"Partially generated using [mypy-protobuf==$MYPY_PROTOBUF_VERSION](https:\/\/github.com\/nipunn1313\/mypy-protobuf\/tree\/v$MYPY_PROTOBUF_VERSION) on tensorflow==$TENSORFLOW_VERSION\"/" \
+    "$REPO_ROOT/stubs/tensorflow/METADATA.toml"
+
+# Cleanup last. If the script fails halfway, it's nice to be able to re-run it immediatly
+rm -rf repository/

stubs/protobuf/METADATA.toml

@@ -1,6 +1,6 @@
 version = "4.24.*"
 upstream_repository = "https://github.com/protocolbuffers/protobuf"
-extra_description = "Generated with aid from [mypy-protobuf](https://github.com/nipunn1313/mypy-protobuf) v3.5.0"
+extra_description = "Generated using [mypy-protobuf==3.5.0](https://github.com/nipunn1313/mypy-protobuf/tree/v3.5.0) on protobuf==4.21.8"
 partial_stub = true
 
 [tool.stubtest]

stubs/tensorflow/METADATA.toml

@@ -2,6 +2,7 @@ version = "2.12.*"
 upstream_repository = "https://github.com/tensorflow/tensorflow"
 # requires a version of numpy with a `py.typed` file
 requires = ["numpy>=1.20", "types-protobuf"]
+extra_description = "Partially generated using [mypy-protobuf==3.5.0](https://github.com/nipunn1313/mypy-protobuf/tree/v3.5.0) on tensorflow==2.12.1"
 partial_stub = true
 
 [tool.stubtest]

stubs/tensorflow/tensorflow/compiler/xla/autotune_results_pb2.pyi

@@ -0,0 +1,94 @@
+"""
+@generated by mypy-protobuf.  Do not edit manually!
+isort:skip_file
+Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+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.
+==============================================================================
+"""
+import builtins
+import collections.abc
+import google.protobuf.descriptor
+import google.protobuf.internal.containers
+import google.protobuf.message
+import sys
+import tensorflow.tsl.protobuf.autotuning_pb2
+
+if sys.version_info >= (3, 8):
+    import typing as typing_extensions
+else:
+    import typing_extensions
+
+DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
+
+@typing_extensions.final
+class AutotuneResults(google.protobuf.message.Message):
+    """A collection of algorithms for particular dot/convs.  Usually this is "the
+    best" algorithm for the particular dot/conv, although that's not strictly
+    required.
+
+    Users don't interact with this proto directly.  It's used internally to
+    facilitate ahead-of-time autotuning -- The string used by
+    xla::{Serialize,Load}AutotuneResults is, internally, a serialization of this
+    proto.
+
+    LINT.IfChange
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    @typing_extensions.final
+    class Entry(google.protobuf.message.Message):
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+        DEVICE_FIELD_NUMBER: builtins.int
+        HLO_FIELD_NUMBER: builtins.int
+        RESULT_FIELD_NUMBER: builtins.int
+        device: builtins.str
+        hlo: builtins.str
+        @property
+        def result(self) -> tensorflow.tsl.protobuf.autotuning_pb2.AutotuneResult:
+            """nb: These results are always tied to a particular version of
+            cublas/cudnn, but this is *especially* true for cublasLt results.  For
+            cublasLt gemms, the result is an index into the list of candidate
+            algorithms returned by cublasLt.  Different version of cublasLt ->
+            different list of algos -> different interpretation of results!
+            """
+        def __init__(
+            self,
+            *,
+            device: builtins.str | None = ...,
+            hlo: builtins.str | None = ...,
+            result: tensorflow.tsl.protobuf.autotuning_pb2.AutotuneResult | None = ...,
+        ) -> None: ...
+        def HasField(self, field_name: typing_extensions.Literal["result", b"result"]) -> builtins.bool: ...
+        def ClearField(self, field_name: typing_extensions.Literal["device", b"device", "hlo", b"hlo", "result", b"result"]) -> None: ...
+
+    VERSION_FIELD_NUMBER: builtins.int
+    DOTS_FIELD_NUMBER: builtins.int
+    CONVS_FIELD_NUMBER: builtins.int
+    version: builtins.int
+    @property
+    def dots(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___AutotuneResults.Entry]: ...
+    @property
+    def convs(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___AutotuneResults.Entry]: ...
+    def __init__(
+        self,
+        *,
+        version: builtins.int | None = ...,
+        dots: collections.abc.Iterable[global___AutotuneResults.Entry] | None = ...,
+        convs: collections.abc.Iterable[global___AutotuneResults.Entry] | None = ...,
+    ) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["convs", b"convs", "dots", b"dots", "version", b"version"]) -> None: ...
+
+global___AutotuneResults = AutotuneResults