feat(python): add editor typing and API documentation

.github/scripts/ci_changes.py

@@ -125,7 +125,7 @@ def classify(paths):
             )
         elif path.startswith("python_bindings/"):
             selected.update(("python_build", "wheels", "benchmarks"))
-            if path.endswith(".py"):
+            if path.endswith((".py", ".pyi")):
                 selected.add("python_quality")
             elif path.endswith(CPP_SUFFIXES):
                 selected.update(("cpp_format", "cpp_tidy"))

.github/scripts/test_ci_workflow_contracts.py

@@ -6,6 +6,8 @@
 import subprocess
 import sys
 import tempfile
+import textwrap
+import tomllib
 import unittest
 from pathlib import Path
 
@@ -15,6 +17,62 @@
 
 
 class WorkflowContractsTest(unittest.TestCase):
+    def test_installed_typing_checks_are_required(self):
+        config = tomllib.loads((ROOT / "pyproject.toml").read_text())
+        test_dependencies = config["project"]["optional-dependencies"]["test"]
+        wheel_dependencies = config["tool"]["cibuildwheel"]["test-requires"]
+        checker = [value for value in test_dependencies if value.startswith("mypy==")]
+        self.assertEqual(len(checker), 1)
+        self.assertIn(checker[0], wheel_dependencies)
+        self.assertFalse(
+            any(value.startswith("mypy") for value in config["project"]["dependencies"])
+        )
+        workflow = (ROOT / ".github/workflows/python.yml").read_text()
+        build = workflow.split("  python-build:\n", 1)[1].split(
+            "  python-result:\n", 1
+        )[0]
+        self.assertIn('pip install ".[test]"', build)
+        self.assertIn("python -m pytest tests/python/test_editor_support.py", build)
+        self.assertNotIn("continue-on-error", build)
+        self.assertIn(
+            "pytest {project}/tests/python",
+            config["tool"]["cibuildwheel"]["test-command"],
+        )
+        self.assertNotIn(
+            "--ignore={project}/tests/python/test_editor_support.py",
+            config["tool"]["cibuildwheel"]["test-command"],
+        )
+        for path in (
+            "python_bindings/__init__.pyi",
+            "tests/python/typing/public_api.py",
+        ):
+            with self.subTest(path=path):
+                self.assertLessEqual(
+                    {"python_build", "python_quality", "wheels"}, classify([path])
+                )
+        # Exercise the actual aggregator script, so a failed selected matrix
+        # cannot be reported as a successful Python validation check.
+        result = workflow.split("  python-result:\n", 1)[1].split(
+            "  python-windows:\n", 1
+        )[0]
+        self.assertIn("needs: [changes, python-build]", result)
+        script = textwrap.dedent(result.split("        run: |\n", 1)[1])
+        for matrix_result in ("success", "failure", "cancelled", "skipped"):
+            with self.subTest(matrix_result=matrix_result):
+                checked = subprocess.run(
+                    ["bash", "-e", "-c", script],
+                    env=dict(
+                        os.environ,
+                        CHANGES_RESULT="success",
+                        PYTHON_SELECTED="true",
+                        MATRIX_RESULT=matrix_result,
+                    ),
+                    capture_output=True,
+                    text=True,
+                    check=False,
+                )
+                self.assertEqual(checked.returncode == 0, matrix_result == "success")
+
     def test_tidy_rejects_empty_scope_and_propagates_analysis_failure(self):
         with tempfile.TemporaryDirectory() as directory:
             work = Path(directory)

.github/workflows/python.yml

@@ -68,6 +68,11 @@ jobs:
       - name: Build and install wheel
         run: pip install ".[test]" -v
 
+      - name: Check installed Python types and editor metadata
+        env:
+          OMP_NUM_THREADS: "2"
+        run: python -m pytest tests/python/test_editor_support.py --tb=short -q
+
       # Repaired-wheel jobs run the full Python suite.
       - name: Smoke-test source installation
         run: >-

pyproject.toml

@@ -27,7 +27,7 @@ Issues = "https://github.com/VectorDB-NTU/RaBitQ-Library/issues"
 Paper = "https://doi.org/10.1145/3725413"
 
 [project.optional-dependencies]
-test = ["pytest>=7", "numpy"]
+test = ["pytest>=7", "numpy", "mypy==2.4.0"]
 
 [tool.scikit-build]
 # Disable directory auto-discovery so we can map it via CMake's install target
@@ -45,7 +45,7 @@ cmake.args = [
 build = "cp311-manylinux_x86_64 cp312-manylinux_x86_64 cp313-manylinux_x86_64 cp314-manylinux_x86_64 cp311-manylinux_aarch64 cp312-manylinux_aarch64 cp313-manylinux_aarch64 cp314-manylinux_aarch64 cp311-win_amd64 cp312-win_amd64 cp313-win_amd64 cp314-win_amd64 cp311-macosx_arm64 cp312-macosx_arm64 cp313-macosx_arm64 cp314-macosx_arm64"
 build-verbosity = 1
 environment = { CMAKE_ARGS = "-DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF" }
-test-requires = ["pytest"]
+test-requires = ["pytest", "mypy==2.4.0"]
 test-command = "pytest {project}/tests/python --ignore={project}/tests/python/test_compare_with_faiss.py --ignore={project}/tests/python/test_cpp_clustering_examples.py --tb=short -q"
 
 [tool.cibuildwheel.linux]

python_bindings/CMakeLists.txt

@@ -58,7 +58,7 @@ endif()
 # 5. Map files to a 'rabitqlib' directory
 # inside the target wheel, completely bypassing your local layout.
 install(TARGETS _rabitqlib DESTINATION rabitqlib)
-install(FILES __init__.py DESTINATION rabitqlib)
+install(FILES __init__.py __init__.pyi py.typed DESTINATION rabitqlib)
 
 if(RABITQ_BUILD_BENCHMARKS)
     # Benchmark the production conversion helper without importing or replacing

python_bindings/__init__.py

@@ -31,10 +31,24 @@ def __init__(self, d, k, **kwargs):
         self._impl = None
 
     def train(self, x):
-        """Train on ``(n, d)`` data and return the last pre-update objective.
-
-        ``final_obj`` measures the returned assignments against the final centroids.
-        A failed fit preserves the results of the previous successful fit.
+        """Train centroids and retain the final assignments.
+
+        Parameters
+        ----------
+        x : array_like of float, shape (n, d)
+            Vectors in the original, unpadded coordinates. Converted to contiguous
+            float32; a copy may be made.
+
+        Returns
+        -------
+        objective : float
+            Last pre-update objective. This can differ from final_obj, which uses
+            returned assignments and final centroids.
+
+        Notes
+        -----
+        A failed fit preserves the last successful results. Returned arrays are
+        copies.
         """
         implementation = self._implementation_type(self.d, self.k, self.cp)
         objective = implementation.train(x)
@@ -43,38 +57,137 @@ def train(self, x):
 
     @property
     def centroids(self):
+        """Final centroids, float32 array of shape (k, d), or None before training."""
         return None if self._impl is None else self._impl.centroids
 
     @property
     def assignments(self):
+        """Final centroid IDs, uint32 array of shape (n,), or None before training."""
         return None if self._impl is None else self._impl.assignments
 
     @property
     def distances(self):
+        """Final assignment distances, float32 array of shape (n,), or None before training."""
         return None if self._impl is None else self._impl.distances
 
     @property
     def obj(self):
+        """Pre-update objectives, float64 array of shape (iterations,), or None before training."""
         return None if self._impl is None else self._impl.obj
 
     @property
     def final_obj(self):
+        """Objective for final assignments and centroids, or None before training."""
         return None if self._impl is None else self._impl.final_obj
 
     @property
     def iteration_stats(self):
+        """Per-iteration statistics, or None before training."""
         return None if self._impl is None else self._impl.iteration_stats
 
 
 class QGKMeans(_KMeans):
-    """SymphonyQG graph assignment; recommended for many centroids."""
+    """Cluster vectors using SymphonyQG graph assignment.
+
+    Parameters
+    ----------
+    d : int
+        Original vector dimension, from 64 through 65536.
+    k : int
+        Number of centroids; greater than graph_degree and no more than the
+        number of training vectors.
+    niter : int, optional
+        Maximum number of clustering iterations. Default is 25.
+    verbose : bool, optional
+        Print training progress when True. Default is False.
+    spherical : bool, optional
+        Normalize centroids after each update and use inner-product assignment.
+        Normalize input vectors yourself for cosine clustering. Default is
+        False.
+    seed : int, optional
+        Random seed. Default is 42.
+    min_points_per_centroid : int, optional
+        Warn below this many training vectors per centroid. Default is 39.
+    early_stop_threshold : float, optional
+        Stop when the absolute relative objective change is at most this
+        threshold. Default is 0.0.
+    num_threads : int, optional
+        Worker count; 0 selects the available CPU count. Default is 0.
+    final_assignment : FinalAssignmentMode, optional
+        Approximate uses the training assignment method; Exact uses exhaustive
+        assignment against final centroids. Default is
+        FinalAssignmentMode.Approximate.
+    graph_degree : int, optional
+        Centroid graph degree; a positive multiple of 32. Default is 32.
+    ef_build : int, optional
+        Positive centroid-graph construction search breadth. Default is 240.
+    ef_search : int, optional
+        Positive centroid-graph assignment search breadth. Default is 16.
+    graph_build_iterations : int, optional
+        Positive number of graph-building passes after initialization. Default
+        is 1.
+    quantization_bits : int, optional
+        Centroid storage: 0 for raw float32, or 4 or 8 for RaBitQ. Default is
+        0.
+
+    Notes
+    -----
+    Training results are None until the first successful train call. Set cp
+    fields before training to customize subsequent runs. A failed train
+    preserves previous results.
+    """
 
     _parameters_type = QGKMeansParameters
     _implementation_type = _QGKMeans
 
 
 class RaBitQKMeans(_KMeans):
-    """Flat RaBitQ assignment; recommended for few centroids."""
+    """Cluster vectors using flat RaBitQ assignment.
+
+    Parameters
+    ----------
+    d : int
+        Original vector dimension, from 64 through 65536.
+    k : int
+        Number of centroids, between 1 and the number of training vectors.
+    niter : int, optional
+        Maximum number of clustering iterations. Default is 25.
+    verbose : bool, optional
+        Print training progress when True. Default is False.
+    spherical : bool, optional
+        Normalize centroids after each update and use inner-product assignment.
+        Normalize input vectors yourself for cosine clustering. Default is
+        False.
+    seed : int, optional
+        Random seed. Default is 42.
+    min_points_per_centroid : int, optional
+        Warn below this many training vectors per centroid. Default is 39.
+    early_stop_threshold : float, optional
+        Stop when the absolute relative objective change is at most this
+        threshold. Default is 0.0.
+    num_threads : int, optional
+        Worker count; 0 selects the available CPU count. Default is 0.
+    final_assignment : FinalAssignmentMode, optional
+        Approximate uses the training assignment method; Exact uses exhaustive
+        assignment against final centroids. Default is
+        FinalAssignmentMode.Approximate.
+
+    Notes
+    -----
+    Training results are None until the first successful train call. Set cp
+    fields before training to customize subsequent runs. A failed train
+    preserves previous results.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from rabitqlib import RaBitQKMeans
+    >>> x = np.random.default_rng(42).normal(size=(64, 64)).astype(np.float32)
+    >>> model = RaBitQKMeans(64, 2, niter=2, num_threads=1)
+    >>> objective = model.train(x)
+    >>> model.centroids.shape
+    (2, 64)
+    """
 
     _parameters_type = RaBitQKMeansParameters
     _implementation_type = _RaBitQKMeans