#!/usr/bin/env bash

# Generate type stubs for .proto definition files.

# This should be run from as
# ./gen_proto_typestubs.sh
# (i.e., from inside the proto/ directory)

# assumes mypy-protobuf installed to ~/.local; i.e. via
# pip3 install mypy-protobuf --user

set -euxo pipefail

MYPY_PROTOBUF_HOME="${1:-${HOME}/.local/bin}"

pushd ../../
buck run fbsource//third-party/protobuf:protoc -- --plugin=protoc-gen-mypy="${MYPY_PROTOBUF_HOME}"/protoc-gen-mypy --mypy_out=./ caffe2/proto/*.proto
popd

# get rid of 'builtins.' prefix, which pyre does not like
sed -E -i 's/builtins\.//g' ./*.pyi

# mypy-protobuf references types from other mypy-protobuf-generated stubs as
# 'type.V', but it should just be 'type', so we get rid of the '.V' suffix
# when it's not followed by parens to indicate a particular enum value.
sed -E -i 's/\.V([^(_[:alnum:]])/\1/g' ./*.pyi

# ---------------------------
# Freedom-patched DeviceTypes
# ---------------------------
#
# In order to make DeviceTypes like CPU, CUDA, etc. directly accessible from
# the caffe2_pb2 module, they are currently freedom-patched into it in
# caffe2/python/__init__.py. This is not ideal: it would be better if these
# were autogenerated when the protobuf definitions were created by using
# allow_alias = true in the DeviceTypeProto definition in caffe2.proto.
#
# However, it is impossible to do this currently without significant effort.
# The issue is that the generated proto constants would conflict with various
# constants defined in the C++ caffe2 codebase (`caffe2_pb2.h`). We cannot
# simply remove these constants and replace them with the caffe2
# DeviceTypeProto constants, because a huge portion of code expects
# at::DeviceType constants defined in `core/DeviceType.h` (apparently
# duplicated to avoid having to figure out how to autogenerate the protobuf
# definitions using cmake for ATen).
#
# Instead, we make a best-effort to add additional definitions in
# `caffe2_pb2.py` by looking for any freedom-patched constants in
# `caffe2/python/__init__.py` and making sure they have corresponding stubs in
# the pyi (see `gen_proto_typestubs_helper.py`).

python3 ./gen_proto_typestubs_helper.py >> caffe2_pb2.pyi
