Add PEP 561 type stubs for chiavdf

.github/workflows/build.yml

@@ -197,15 +197,28 @@ jobs:
       with:
         python-version: '3.12'
 
+    - name: Install build dependencies
+      run: |
+        sudo apt-get update
+        sudo apt-get install -y cmake libgmp-dev libgmpxx4ldbl
+
     - name: flake8
       run: |
         pip install flake8
-        flake8 tests setup.py
+        flake8 tests setup.py chiavdf
 
     - name: mypy
       run: |
         pip install mypy
-        mypy --config-file mypi.ini setup.py tests
+        mypy --config-file mypi.ini setup.py tests chiavdf
+
+    - name: stubtest
+      env:
+        BUILD_VDF_CLIENT: N
+      run: |
+        pip install mypy
+        pip install .
+        stubtest --concise chiavdf
 
   upload:
     name: Upload to PyPI - Ubuntu 3.12 Intel

chiavdf/__init__.py

@@ -0,0 +1,21 @@
+from chiavdf._chiavdf import (
+    bqfc_deserialize,
+    create_discriminant,
+    create_discriminant_and_verify_n_wesolowski,
+    get_b_from_n_wesolowski,
+    prove,
+    verify_n_wesolowski,
+    verify_n_wesolowski_with_b,
+    verify_wesolowski,
+)
+
+__all__ = [
+    "bqfc_deserialize",
+    "create_discriminant",
+    "create_discriminant_and_verify_n_wesolowski",
+    "get_b_from_n_wesolowski",
+    "prove",
+    "verify_n_wesolowski",
+    "verify_n_wesolowski_with_b",
+    "verify_wesolowski",
+]

chiavdf/__init__.pyi

@@ -0,0 +1,62 @@
+__all__ = [
+    "bqfc_deserialize",
+    "create_discriminant",
+    "create_discriminant_and_verify_n_wesolowski",
+    "get_b_from_n_wesolowski",
+    "prove",
+    "verify_n_wesolowski",
+    "verify_n_wesolowski_with_b",
+    "verify_wesolowski",
+]
+
+def create_discriminant(challenge_hash: bytes, discriminant_size_bits: int) -> str: ...
+def verify_wesolowski(
+    discriminant: str,
+    x_s: bytes | str,
+    y_s: bytes | str,
+    proof_s: bytes | str,
+    num_iterations: int,
+) -> bool: ...
+def verify_n_wesolowski(
+    discriminant: str,
+    x_s: bytes | str,
+    proof_blob: bytes,
+    num_iterations: int,
+    disc_size_bits: int,
+    recursion: int,
+) -> bool: ...
+def create_discriminant_and_verify_n_wesolowski(
+    challenge_hash: bytes,
+    discriminant_size_bits: int,
+    x_s: bytes | str,
+    proof_blob: bytes,
+    num_iterations: int,
+    recursion: int,
+) -> bool: ...
+def prove(
+    challenge_hash: bytes,
+    x_s: bytes | str,
+    discriminant_size_bits: int,
+    num_iterations: int,
+    shutdown_file_path: str,
+) -> bytes: ...
+def verify_n_wesolowski_with_b(
+    discriminant: str,
+    B: str,
+    x_s: bytes | str,
+    proof_blob: bytes,
+    num_iterations: int,
+    recursion: int,
+) -> tuple[bool, bytes]: ...
+def bqfc_deserialize(
+    discriminant: str,
+    data: bytes,
+    strict: bool = True,
+) -> tuple[bytes, bytes]: ...
+def get_b_from_n_wesolowski(
+    discriminant: str,
+    x_s: bytes | str,
+    proof_blob: bytes,
+    num_iterations: int,
+    recursion: int,
+) -> str: ...

pyproject.toml

@@ -8,7 +8,7 @@ local_scheme = "no-local-version"
 
 [tool.cibuildwheel]
 test-requires = "pytest"
-test-command = "pytest -v {project}/tests"
+test-command = "mv {project}/chiavdf {project}/DO_NOT_IMPORT && pytest -v {project}/tests"
 skip = "cp31?t-* *-manylinux_i686 *-win32 *-musllinux_*"
 
 [tool.cibuildwheel.linux]

setup.py

@@ -117,7 +117,9 @@ def build_extension(self, ext):
     long_description=_long_description,
     long_description_content_type="text/markdown",
     url="https://github.com/Chia-Network/chiavdf",
-    ext_modules=[CMakeExtension("chiavdf", "src")],
+    packages=["chiavdf"],
+    package_data={"chiavdf": ["py.typed", "__init__.pyi"]},
+    ext_modules=[CMakeExtension("chiavdf._chiavdf", "src")],
     cmdclass=dict(build_ext=CMakeBuild, build_hook=build_hook),
     zip_safe=False,
 )

src/CMakeLists.txt

@@ -246,19 +246,19 @@ if(BUILD_PYTHON)
     FetchContent_MakeAvailable(pybind11-src)
   endif()
 
-  pybind11_add_module(chiavdf
+  pybind11_add_module(_chiavdf
     ${CMAKE_CURRENT_SOURCE_DIR}/python_bindings/fastvdf.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/refcode/lzcnt.c
   )
 
-  target_link_libraries(chiavdf PRIVATE ${GMP_LIBRARIES} ${GMPXX_LIBRARIES})
+  target_link_libraries(_chiavdf PRIVATE ${GMP_LIBRARIES} ${GMPXX_LIBRARIES})
   if(UNIX)
-    target_link_libraries(chiavdf PRIVATE -pthread)
+    target_link_libraries(_chiavdf PRIVATE -pthread)
   endif()
   if (WIN32)
     # workaround for constexpr mutex constructor change in MSVC 2022
     # https://stackoverflow.com/questions/78598141/first-stdmutexlock-crashes-in-application-built-with-latest-visual-studio
-    target_compile_definitions(chiavdf PUBLIC _DISABLE_CONSTEXPR_MUTEX_CONSTRUCTOR)
+    target_compile_definitions(_chiavdf PUBLIC _DISABLE_CONSTEXPR_MUTEX_CONSTRUCTOR)
   endif()
 endif()
 

src/python_bindings/fastvdf.cpp

@@ -25,7 +25,7 @@ static py::bytes to_signed_bytes_be(const integer& value) {
     return py::bytes(out);
 }
 
-PYBIND11_MODULE(chiavdf, m) {
+PYBIND11_MODULE(_chiavdf, m) {
     m.doc() = "Chia proof of time";
 
     // Creates discriminant.