SNOW-1791994 IR: Generate typed protobuf Python code

.github/scripts/install_protoc.sh

@@ -71,4 +71,8 @@ install() {
 
 install
 
+# mypy-protobuf is used to generated typed Python code from protobuf
+pip install mypy-protobuf
+echo "mypy-protobuf version: $(protoc-gen-mypy --version)"
+
 echo "${SCRIPT_NAME} done."

.github/workflows/precommit.yml

@@ -593,19 +593,13 @@ jobs:
           python-version: '3.9'
       - name: Upgrade setuptools and pip
         run: python -m pip install -U setuptools pip
+      - name: Install tox
+        run: python -m pip install tox
       - name: Install protoc
         shell: bash
         run: .github/scripts/install_protoc.sh
-      - name: Install Snowpark
-        run: python -m pip install ".[modin-development]"
-      - name: Install Sphinx
-        run: python -m pip install sphinx
       - name: Build document
-        working-directory: docs
-        # treat warning as failure but complete the entire process
-        run: |
-          make clean
-          make html SPHINXOPTS="-W --keep-going"
+        run: python -m tox -e docs
       - name: Upload html files
         uses: actions/upload-artifact@v4
         with:

.pre-commit-config.yaml

@@ -91,6 +91,7 @@ repos:
       - id: mypy
         files: >
           (?x)^(
+            src/snowflake/snowpark/_internal/ast/.*\.py|
             src/snowflake/snowpark/modin/pandas/indexing.py|
             src/snowflake/snowpark/modin/plugin/_internal/.*\.py|
             src/snowflake/snowpark/modin/pandas/snow_partition_iterator.py|
@@ -103,9 +104,12 @@ repos:
           - --disallow-incomplete-defs
         additional_dependencies:
           - types-requests
+          - types-python-dateutil
           - types-urllib3
           - types-setuptools
           - types-pyOpenSSL
           - types-setuptools
+          - mypy-protobuf
+          - typed-ast
           - pytest
           - numpy

scripts/jenkins_regress.sh

@@ -18,7 +18,7 @@ exit_code_decorator(){
 gpg --quiet --batch --yes --decrypt --passphrase="$GPG_KEY" --output "tests/parameters.py" scripts/parameters.py.gpg
 
 # Install protoc
-pip install protoc-wheel-0==21.1
+pip install protoc-wheel-0==21.1 mypy-protobuf
 
 # Run linter, Python 3.8 test and code coverage jobs
 exit_code_decorator "python -m tox -c $WORKING_DIR" -e notdoctest

scripts/jenkins_regress_snowpandas.sh

@@ -15,7 +15,7 @@ gpg --quiet --batch --yes --decrypt --passphrase="$GPG_KEY" --output "tests/para
 python -m pip install tox
 
 # Install protoc
-pip install protoc-wheel-0==21.1
+pip install protoc-wheel-0==21.1 mypy-protobuf
 
 # Run snowpandas tests
 python -m tox -c $WORKING_DIR -e snowparkpandasjenkins-modin

setup.py

@@ -57,6 +57,7 @@
     "pytest-assume",  # sql counter check
     "decorator",  # sql counter check
     "protoc-wheel-0==21.1",  # Protocol buffer compiler, for Snowpark IR
+    "mypy-protobuf",  # used in generating typed Python code from protobuf for Snowpark IR
     "lxml",  # used in read_xml tests
 ]
 
@@ -81,7 +82,7 @@
 
 if protoc is None:
     sys.stderr.write(
-        "protoc is not installed nor found. Please install the binary package, e.g., `pip install protoc-wheel-0==21.1`\n"
+        "protoc is not installed nor found. Please install the binary package, e.g., `pip install protoc-wheel-0==21.1 mypy-protobuf`\n"
     )
     sys.exit(-1)
 
@@ -115,7 +116,7 @@ def generate_proto(source):
             protoc,
             f"--proto_path={proto_dir}",
             f"--python_out={output_dir}",
-            f"--pyi_out={output_dir}",
+            f"--mypy_out={output_dir}",
             source,
         ]
         if subprocess.call(protoc_command) != 0:

src/snowflake/snowpark/_internal/ast/batch.py

@@ -64,9 +64,9 @@ def __init__(self, session: Session) -> None:
         self._init_batch()
 
         # Track callables in this dict (memory id -> TrackedCallable).
-        self._callables = {}
+        self._callables: dict[int, TrackedCallable] = {}
 
-    def reset_id_gen(self):
+    def reset_id_gen(self) -> None:
         """Resets the ID generator."""
         self._id_gen = itertools.count(start=1)
 
@@ -84,7 +84,7 @@ def assign(self, symbol: Optional[str] = None) -> proto.Assign:
         stmt.assign.symbol.value = symbol if isinstance(symbol, str) else ""
         return stmt.assign
 
-    def eval(self, target: proto.Assign):
+    def eval(self, target: proto.Assign) -> None:
         """
         Creates a new evaluation statement.
 
@@ -102,7 +102,7 @@ def flush(self) -> SerializedBatch:
         self._init_batch()
         return SerializedBatch(req_id, batch)
 
-    def _init_batch(self):
+    def _init_batch(self) -> None:
         # Reset the AST batch by initializing a new request.
         self._request_id = AstBatch.generate_request_id()  # Generate a new unique ID.
         self._request = proto.Request()