Add Python type stub (.pyi) generation for tesseract_decoder (#189)

Addresses #154.

1. Creates python stub file generation script using pybind11-stubgen.
This script is then called via a genrule when building the Wheel

2. Adds tests to ensure stubfiles exist and are valid python (maybe we
can drop the valid python check + symbol expectation, WDYT?). Note that
I think there is an issue with the visualizer stubs since

Smaller changes:
a. Adds pybind11-stubgen to the requirements file.


Verification:
Unit tests, and just seeing if autocomplete works now:
Autocomplete without stubfiles:
<img width="396" height="161" alt="Screenshot 2026-02-20 at 9 12 48 AM"
src="https://github.com/user-attachments/assets/c2697630-e9f0-4041-9ddb-fe4f60558b0a"
/>

Autocomplete with stubfiles:
<img width="740" height="355" alt="Screenshot 2026-02-20 at 9 09 29 AM"
src="https://github.com/user-attachments/assets/48843008-2558-4379-ab65-d6605fc4c3af"
/>

Author: dandragona-dev

BUILD

@@ -20,6 +20,7 @@ py_wheel(
     distribution = "tesseract_decoder",
     deps=[
         "//src:tesseract_decoder",
+        "//src/py:generated_stubs",
         "//src/py/_tesseract_py_util:_tesseract_py_util",
         ":package_data",
     ],

src/BUILD

@@ -72,6 +72,7 @@ pybind_library(
         "tesseract.pybind.h",
         "tesseract_sinter_compat.pybind.h",
     ],
+    copts = OPT_COPTS,
     deps = [
         ":libcommon",
         ":libutils",
@@ -85,6 +86,7 @@ pybind_extension(
     srcs = [
         "tesseract.pybind.cc",
     ],
+    copts = OPT_COPTS,
     deps = [
         ":tesseract_decoder_pybind",
     ],

src/py/BUILD

@@ -15,6 +15,7 @@
 load("@rules_python//python:py_test.bzl", "py_test")
 load("@rules_python//python:pip.bzl", "compile_pip_requirements")
 load("@rules_python//python:py_library.bzl", "py_library")
+load("@rules_python//python:py_binary.bzl", "py_binary")
 
 py_library(
     name = "shared_decoding_tests",
@@ -94,6 +95,46 @@ py_test(
 
 
 
+py_test(
+    name = "stub_test",
+    srcs = ["stub_test.py"],
+    data = [":generated_stubs"],
+    visibility = ["//:__subpackages__"],
+    deps = [
+        "@pypi//pytest",
+    ],
+    imports = ["..", "."],
+)
+
+py_binary(
+    name = "generate_stubs",
+    srcs = ["generate_stubs.py"],
+    deps = [
+        "//src:lib_tesseract_decoder",
+        "@pypi//pybind11_stubgen",
+        "@pypi//stim",
+    ],
+    imports = ["..", "."],
+)
+
+STUB_FILES = [
+    "__init__.pyi",
+    "common.pyi",
+    "simplex.pyi",
+    "tesseract.pyi",
+    "tesseract_sinter_compat.pyi",
+    "utils.pyi",
+    "viz.pyi",
+]
+
+genrule(
+    name = "generated_stubs",
+    tools = [":generate_stubs"],
+    outs = ["tesseract_decoder-stubs/" + f for f in STUB_FILES],
+    cmd = "$(location :generate_stubs) --output-dir $(@D)",
+    visibility = ["//visibility:public"],
+)
+
 compile_pip_requirements(
     name = "requirements",
     src = "requirements.in",

src/py/generate_stubs.py

@@ -0,0 +1,142 @@
+#!/usr/bin/env python3
+# Copyright 2025 Google LLC
+#
+# 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.
+
+"""Generate Python type stub (.pyi) files for the tesseract_decoder module.
+
+This script uses pybind11-stubgen to produce .pyi stub files for the tesseract decoder
+module's C++ API. The .pyi stub files provide information for IDEs.
+
+Built from //src/py:generated_stubs which uses a genrule to call this script.
+"""
+
+import argparse
+import os
+import sys
+
+
+def _ensure_module_importable():
+    """Ensure tesseract_decoder is importable, adjusting sys.path if needed."""
+
+    try:
+        import tesseract_decoder
+        return tesseract_decoder
+    except Exception as e:
+        import traceback
+        print("====== DEBUG INFO ======", file=sys.stderr)
+        print(f"Exception: {type(e).__name__}: {e}", file=sys.stderr)
+        print(f"sys.path: {sys.path}", file=sys.stderr)
+        print(f"__file__: {__file__}", file=sys.stderr)
+        print(f"cwd: {os.getcwd()}", file=sys.stderr)
+        print("========================", file=sys.stderr)
+        print(
+            f"ERROR: Cannot import tesseract_decoder.\n"
+            "Ensure the compiled module is on sys.path or PYTHONPATH.",
+            file=sys.stderr,
+        )
+        traceback.print_exc()
+        sys.exit(1)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Generate .pyi stubs for tesseract_decoder"
+    )
+    parser.add_argument(
+        "--output-dir",
+        default=None,
+        help="Directory to place the generated .pyi file. "
+        "Defaults to the directory containing the tesseract_decoder module.",
+    )
+    args = parser.parse_args()
+
+    module = _ensure_module_importable()
+    module_file = os.path.abspath(module.__file__)
+    module_dir = os.path.dirname(module_file)
+
+    output_dir = os.path.abspath(args.output_dir) if args.output_dir else module_dir
+
+    print(f"Generating stubs for tesseract_decoder...")
+    print(f"  Module location: {module_file}")
+    print(f"  Output directory: {output_dir}")
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Use pybind11_stubgen programmatically.
+    try:
+        from pybind11_stubgen import main as stubgen_main
+    except ImportError as e:
+        print(f"ImportError: {e}", file=sys.stderr)
+        print(f"sys.path: {sys.path}", file=sys.stderr)
+        print(
+            "ERROR: pybind11-stubgen is not installed. "
+            "Install with: pip install pybind11-stubgen",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
+    # Build argv for pybind11-stubgen CLI.
+    # --enum-class-locations maps enum names to their fully-qualified module path
+    # so pybind11-stubgen can resolve default values like <DetOrder.DetBFS: 0>.
+    stubgen_argv = [
+        "pybind11-stubgen",
+        "tesseract_decoder",
+        "--output-dir",
+        output_dir,
+        "--enum-class-locations",
+        "DetOrder:tesseract_decoder.utils",
+    ]
+
+    # Save and restore sys.argv since pybind11-stubgen uses argparse.
+    old_argv = sys.argv
+    sys.argv = stubgen_argv
+    try:
+        stubgen_main()
+    except SystemExit as e:
+        if e.code != 0:
+            print(f"ERROR: pybind11-stubgen exited with code {e.code}", file=sys.stderr)
+            sys.exit(1)
+    finally:
+        sys.argv = old_argv
+
+    # Verify the output exists, and append -stubs to the directory name.
+    original_stub_pkg_dir = os.path.join(output_dir, "tesseract_decoder")
+    stub_pkg_dir = os.path.join(output_dir, "tesseract_decoder-stubs")
+    if os.path.exists(original_stub_pkg_dir):
+        import shutil
+        if os.path.exists(stub_pkg_dir):
+            shutil.rmtree(stub_pkg_dir)
+        os.rename(original_stub_pkg_dir, stub_pkg_dir)
+
+    stub_init = os.path.join(stub_pkg_dir, "__init__.pyi")
+
+    if os.path.isfile(stub_init):
+        print(f"Stubs generated successfully at: {stub_pkg_dir}/")
+        for root, dirs, files in os.walk(stub_pkg_dir):
+            for f in sorted(files):
+                if f.endswith(".pyi"):
+                    rel = os.path.relpath(os.path.join(root, f), output_dir)
+                    print(f"  {rel}")
+    else:
+        flat_stub = os.path.join(output_dir, "tesseract_decoder.pyi")
+        if os.path.isfile(flat_stub):
+            print(f"Stubs generated successfully: {flat_stub}")
+        else:
+            print("WARNING: Could not verify stub output location.", file=sys.stderr)
+
+    print("Done.")
+
+
+if __name__ == "__main__":
+    main()

src/py/requirements.in

@@ -1,3 +1,4 @@
 stim
 pytest
-sinter
+sinter
+pybind11-stubgen

src/py/requirements_lock.txt

@@ -455,6 +455,10 @@ pluggy==1.6.0 \
     --hash=sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3 \
     --hash=sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746
     # via pytest
+pybind11-stubgen==2.5.5 \
+    --hash=sha256:10824cd2fc5cbbee032b8fb39e6f6c08de232deb309bc66d786a6c6e8a4601bd \
+    --hash=sha256:758d6d6bbeefc62ad7f78d5e5bbf357ccf6af83cd4504f5f549403f452942708
+    # via -r src/py/requirements.in
 pygments==2.19.1 \
     --hash=sha256:61c16d2a8576dc0649d9f39e089b5f02bcd27fba10d8fb4dcc28173f7a45151f \
     --hash=sha256:9ea1544ad55cecf4b8242fab6dd35a93bbce657034b0611ee383099054ab6d8c

src/py/stub_test.py

@@ -0,0 +1,147 @@
+# Copyright 2025 Google LLC
+#
+# 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.
+
+"""Tests that validate the pre-generated .pyi stub files for tesseract_decoder.
+
+These stubs are generated by `bazel run //src/py:generate_stubs -- --output-dir src`
+and committed to the repository under `src/tesseract_decoder-stubs/`. This test
+validates that the stubs exist, are syntactically valid Python, and contain
+the expected symbols.
+"""
+
+import ast
+import glob
+import os
+
+import pytest
+
+
+def _find_stub_files():
+    """Find all .pyi stub files in the data runfiles."""
+    # Find the src/py/tesseract_decoder-stubs/*.pyi files in the Bazel tree.
+    pattern_genrule = os.path.join(
+            os.environ["TEST_SRCDIR"],
+            os.environ["TEST_WORKSPACE"],
+            "src",
+            "py",
+            "tesseract_decoder-stubs",
+            "*.pyi",
+        )
+    files = glob.glob(pattern_genrule)
+    assert files, f"No stub files found in {pattern_genrule}"
+    return files
+
+
+def _collect_all_names(pyi_files):
+    """Collect all defined names from a list of .pyi files."""
+    all_names = set()
+    for stub_path in pyi_files:
+        with open(stub_path, "r") as f:
+            content = f.read()
+        tree = ast.parse(content)
+        for node in ast.walk(tree):
+            if isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):
+                all_names.add(node.name)
+            elif isinstance(node, ast.Assign):
+                for target in node.targets:
+                    if isinstance(target, ast.Name):
+                        all_names.add(target.id)
+            elif isinstance(node, ast.ImportFrom):
+                if node.names:
+                    for alias in node.names:
+                        all_names.add(
+                            alias.name if alias.asname is None else alias.asname
+                        )
+    return all_names
+
+
+@pytest.fixture(scope="session")
+def stub_files():
+    """Collect all generated .pyi stub files."""
+    files = _find_stub_files()
+    if not files:
+        pytest.skip(
+            "No .pyi stub files found. Run "
+            "'bazel run //src/py:generate_stubs -- --output-dir src' first."
+        )
+    return files
+
+
+class TestStubFilesExist:
+    """Tests that stub files exist and are valid Python."""
+
+    def test_stubs_generated(self, stub_files):
+        """At least one .pyi stub file should exist."""
+        assert len(stub_files) > 0
+
+    EXPECTED_STUBS = [
+        "__init__.pyi",
+        "common.pyi",
+        "simplex.pyi",
+        "tesseract.pyi",
+        "tesseract_sinter_compat.pyi",
+        "utils.pyi",
+        "viz.pyi",
+    ]
+
+    @pytest.mark.parametrize("filename", EXPECTED_STUBS)
+    def test_expected_stub_exists(self, stub_files, filename):
+        """Each expected submodule stub file should be generated."""
+        basenames = [os.path.basename(f) for f in stub_files]
+        assert filename in basenames, (
+            f"Missing expected stub file: {filename}. "
+            f"Found: {basenames}"
+        )
+
+    def test_stubs_are_valid_python(self, stub_files):
+        """All .pyi files should be parseable as valid Python."""
+        for stub_path in stub_files:
+            with open(stub_path, "r") as f:
+                content = f.read()
+            try:
+                ast.parse(content)
+            except SyntaxError as e:
+                basename = os.path.basename(stub_path)
+                pytest.fail(f"Stub file {basename} has invalid syntax: {e}")
+
+
+class TestStubContents:
+    """Tests that the generated stubs contain the expected symbols."""
+
+    EXPECTED_SYMBOLS = [
+        "Symptom",
+        "Error",
+        "TesseractConfig",
+        "TesseractDecoder",
+        "TesseractSinterCompiledDecoder",
+        "TesseractSinterDecoder",
+        "SimplexConfig",
+        "SimplexDecoder",
+        "DetOrder",
+        "Visualizer",
+        "make_tesseract_sinter_decoders_dict",
+    ]
+
+    @pytest.mark.parametrize("symbol", EXPECTED_SYMBOLS)
+    def test_expected_symbol_in_stubs(self, stub_files, symbol):
+        """Key symbols from the pybind11 module should appear in stubs."""
+        all_names = _collect_all_names(stub_files)
+        assert symbol in all_names, (
+            f"Expected symbol '{symbol}' not found in stub files. "
+            f"Found names: {sorted(all_names)}"
+        )
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__]))