Documentation: How can I make pytype notice .pyi stub files without merge-pyi-ing them?

[Byvire]

Greetings from xoogler land. I love pytype and I love protobuf, but I'm having a hard time using them together. Title question seems like the simplest fix to a complex problem, without the pytype and protobuf teams pointing fingers at each other.

TL;DR:

• Pytype's public documentation mentions .pyi stub files but doesn't say how to make pytype not ignore them.
• merge-pyi does not play nice with protobuf generated code. Especially with enums.
• Is there a way to use the pyi files without using merge-pyi?

I put together a demo of the problem I'm having, in https://github.com/Byvire/pytype_proto_bug. See the readme there, and the comments in the scripts build_proto and build_and_merge_proto, which I'll repeat here:

### build_proto
#!/bin/sh

# This version just builds the _pb2.py and _pb2.pyi files and writes them to the
# pytype_proto_demo/proto directory without merging them. After you run this,
# ./run_pytype will fail with errors but ./run_unit_tests will succeed.

protoc -I=. --python_out=. --pyi_out=. pytype_proto_demo/proto/*.proto

### build_and_merge_proto
#!/bin/sh

# This script builds the proto _pb2.py and _pb2.pyi files and then merges them.
# After you run this, you'll find that ./run_unit_tests fails but ./run_pytype
# succeeds.

protoc -I=. --python_out=. --pyi_out=. pytype_proto_demo/proto/*.proto

for pb2_file in $(ls pytype_proto_demo/proto/*_pb2.py) ; do
    if [[ -f "${pb2_file}"i ]] ; then
        merge-pyi -i "${pb2_file}" "${pb2_file}"i
        rm "${pb2_file}"i
    fi
done

If I build the .py and .pyi files without merging them, the unit test passes but pytype gets mad because it doesn't see the .pyi file:

ninja: Entering directory `.pytype'
[1/1] check pytype_proto_demo.example
FAILED: /home/ofk/repo/pytype_proto_bug/.pytype/pyi/pytype_proto_demo/example.pyi 
/home/ofk/.config/.virtualenvs/pytype_bug/bin/python -m pytype.main --imports_info /home/ofk/repo/pytype_proto_bug/.pytype/imports/pytype_proto_demo.example.imports --module-name pytype_proto_demo.example --platform linux -V 3.11 -o /home/ofk/repo/pytype_proto_bug/.pytype/pyi/pytype_proto_demo/example.pyi --analyze-annotated --nofail --quick /home/ofk/repo/pytype_proto_bug/pytype_proto_demo/example.py
File "/home/ofk/repo/pytype_proto_bug/pytype_proto_demo/example.py", line 5, in <module>: No attribute 'Dessert' on module 'pytype_proto_demo.proto.bugdemo_pb2' [module-attr]
File "/home/ofk/repo/pytype_proto_bug/pytype_proto_demo/example.py", line 6, in evaluate: No attribute 'Dessert' on module 'pytype_proto_demo.proto.bugdemo_pb2' [module-attr]
File "/home/ofk/repo/pytype_proto_bug/pytype_proto_demo/example.py", line 8, in evaluate: No attribute 'Dessert' on module 'pytype_proto_demo.proto.bugdemo_pb2' [module-attr]

For more details, see https://google.github.io/pytype/errors.html#module-attr
ninja: build stopped: subcommand failed.
Computing dependencies
Analyzing 1 sources with 1 local dependencies
Leaving directory '.pytype'

On the other hand, if build and merge the files (with the other build script), pytype passes but the unit tests fail:

E
======================================================================
ERROR: pytype_proto_demo.test_example (unittest.loader._FailedTest.pytype_proto_demo.test_example)
----------------------------------------------------------------------
ImportError: Failed to import test module: pytype_proto_demo.test_example
Traceback (most recent call last):
  File "/usr/lib/python3.11/unittest/loader.py", line 419, in _find_test_path
    module = self._get_module_from_name(name)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.11/unittest/loader.py", line 362, in _get_module_from_name
    __import__(name)
  File "/home/ofk/repo/pytype_proto_bug/pytype_proto_demo/test_example.py", line 3, in <module>
    from pytype_proto_demo.proto import bugdemo_pb2
  File "/home/ofk/repo/pytype_proto_bug/pytype_proto_demo/proto/bugdemo_pb2.py", line 13, in <module>
    class Dessert(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
                                 ^^^^^^^^^^^^^^^^^^
NameError: name '_enum_type_wrapper' is not defined


----------------------------------------------------------------------
Ran 1 test in 0.000s

FAILED (errors=1)

The direct cause of that error is that merge-pyi messes up the imports in the merged file. But if I prepend the missing import line

from google.protobuf.internal import enum_type_wrapper as _enum_type_wrapper

then I get a more pernicious error from the same line because

class Dessert(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):

is not allowed in python. You can't make a subclass with a different metaclass than its parent class. So anyway, I thought instead of trying to convince the protobuf team to not do... whatever that is... I'd ask you folks if there's an easy workaround?

(My backup plan is to have two build scripts and run one for unit testing and the other for typechecking, but having code that actually works would be better.)

[Byvire]

Oh yeah, versions of things.

Python 3.11.6.

Pip freeze:

attrs==23.2.0
importlab==0.8.1
Jinja2==3.1.3
libcst==1.2.0
MarkupSafe==2.1.5
mypy-extensions==1.0.0
networkx==3.1
ninja==1.11.1.1
protobuf==4.25.3
pycnite==2023.10.11
pydot==2.0.0
pyparsing==3.1.1
pytype==2024.2.13
-e git+ssh://git@github.com/Byvire/pytype_proto_bug.git@6833fbebf8a4c606ff93b7e3788d1d93cceca54c#egg=pytype_proto_demo
PyYAML==6.0.1
tabulate==0.9.0
toml==0.10.2
typing-inspect==0.9.0
typing_extensions==4.9.0

[Byvire]

Ugh. And libprotoc 25.2