Use pybind11-stubgen to create typestubs for extensions

.github/workflows/cmake.yml

@@ -31,8 +31,21 @@ jobs:
             libbz2-dev \
             ninja-build \
             software-properties-common
+      - name: Install Python 3.12 with the latest version of uv.
+        uses: astral-sh/setup-uv@v7
+        with:
+          python-version: '3.12'
       - uses: actions/checkout@v4
+      - name: Create venv and install build dependencies
+        run: |
+          uv venv --isolated
+          source .venv/bin/activate
+          uv pip install scikit-build-core~=0.10 numpy pybind11-stubgen pybind11~=2.2 setuptools-scm~=9.2.2
       - name: Check CMake configure
-        run: cmake -B build -G Ninja
+        run: |
+          source .venv/bin/activate
+          cmake -B build -G Ninja
       - name: CMake build
-        run: cmake --build build
+        run: |
+          source .venv/bin/activate
+          cmake --build build

CMakeLists.txt

@@ -57,6 +57,8 @@ message(STATUS "HACKDIR set to: ${HACKDIR}")
 set(VARDIR ${HACKDIR})
 set(INSTDIR ${HACKDIR})
 
+# Use FindPython in PyBind11 to find Python.
+set(PYBIND11_FINDPYTHON NEW)
 # pybind11 via FetchContent
 include(FetchContent)
 FetchContent_Declare(
@@ -65,6 +67,7 @@ FetchContent_Declare(
   GIT_TAG v3.0.1
   EXCLUDE_FROM_ALL)
 FetchContent_MakeAvailable(pybind11)
+message(STATUS "Python_EXECUTABLE is set to: ${Python_EXECUTABLE}")
 
 # de-boost-ified version of boost.context via FetchContent
 FetchContent_Declare(
@@ -210,15 +213,52 @@ set_target_properties(_pyconverter PROPERTIES CXX_STANDARD 14)
 target_include_directories(
   _pyconverter PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/third_party/converter)
 
+set(STUB_DIR ${CMAKE_CURRENT_BINARY_DIR}/stubs)
+set_directory_properties(PROPERTIES ADDITIONAL_CLEAN_FILES "${STUB_DIR}")
+add_custom_command(
+  OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/stubs/nle/_pyconverter.pyi
+  COMMAND ${Python_EXECUTABLE} -m pybind11_stubgen _pyconverter --output-dir
+          ${STUB_DIR}
+  WORKING_DIRECTORY $<TARGET_FILE_DIR:_pyconverter>
+  DEPENDS _pyconverter
+  COMMENT "Generating type stubs ..."
+  VERBATIM)
+
+add_custom_command(
+  OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/stubs/nle/_pynethack/__init__.pyi
+  COMMAND ${Python_EXECUTABLE} -m pybind11_stubgen _pynethack --output-dir
+          ${STUB_DIR} --ignore-unresolved-names game_end_types
+  WORKING_DIRECTORY $<TARGET_FILE_DIR:_pynethack>
+  DEPENDS _pynethack
+  COMMENT "Generating type stubs ..."
+  VERBATIM)
+
+add_custom_target(
+  my_module_stubs ALL
+  DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/stubs/nle/_pynethack/__init__.pyi
+          ${CMAKE_CURRENT_BINARY_DIR}/stubs/nle/_pyconverter.pyi)
+
+file(TOUCH ${CMAKE_CURRENT_BINARY_DIR}/py.typed)
+
 set(TILE_FILES "win/share/monsters.txt" "win/share/objects.txt"
                "win/share/other.txt")
 
 install(FILES ${TILE_FILES} DESTINATION ${INSTDIR}/tiles)
+
 # Only install if we are building as part of a Python project.
 if(DEFINED SKBUILD_PROJECT_VERSION)
   install(
     TARGETS _pynethack _pyconverter nethack
     RUNTIME DESTINATION ${PYTHON_PACKAGE_NAME}
     LIBRARY DESTINATION ${PYTHON_PACKAGE_NAME}
     ARCHIVE DESTINATION ${PYTHON_PACKAGE_NAME})
+
+  install(FILES ${CMAKE_CURRENT_BINARY_DIR}/py.typed
+          DESTINATION ${PYTHON_PACKAGE_NAME})
+
+  install(
+    DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/stubs/"
+    DESTINATION ${PYTHON_PACKAGE_NAME}
+    FILES_MATCHING
+    PATTERN "*.pyi")
 endif()

nle/nethack/__init__.py

@@ -1,5 +1,163 @@
 # Copyright (c) Facebook, Inc. and its affiliates.
-from nle._pynethack.nethack import *  # noqa: F403
+import nle._pynethack.nethack
+from nle._pynethack.nethack import AMULET_CLASS
+from nle._pynethack.nethack import ARMOR_CLASS
+from nle._pynethack.nethack import ASCENDED
+from nle._pynethack.nethack import BALL_CLASS
+from nle._pynethack.nethack import BL_MASK_BITS
+from nle._pynethack.nethack import BL_MASK_BLIND
+from nle._pynethack.nethack import BL_MASK_CONF
+from nle._pynethack.nethack import BL_MASK_DEAF
+from nle._pynethack.nethack import BL_MASK_FLY
+from nle._pynethack.nethack import BL_MASK_FOODPOIS
+from nle._pynethack.nethack import BL_MASK_HALLU
+from nle._pynethack.nethack import BL_MASK_LEV
+from nle._pynethack.nethack import BL_MASK_RIDE
+from nle._pynethack.nethack import BL_MASK_SLIME
+from nle._pynethack.nethack import BL_MASK_STONE
+from nle._pynethack.nethack import BL_MASK_STRNGL
+from nle._pynethack.nethack import BL_MASK_STUN
+from nle._pynethack.nethack import BL_MASK_TERMILL
+from nle._pynethack.nethack import BURNING
+from nle._pynethack.nethack import CHAIN_CLASS
+from nle._pynethack.nethack import CHOKING
+from nle._pynethack.nethack import COIN_CLASS
+from nle._pynethack.nethack import COLNO
+from nle._pynethack.nethack import CRUSHING
+from nle._pynethack.nethack import DIED
+from nle._pynethack.nethack import DISSOLVED
+from nle._pynethack.nethack import DROWNING
+from nle._pynethack.nethack import ESCAPED
+from nle._pynethack.nethack import EXPL_MAX
+from nle._pynethack.nethack import FOOD_CLASS
+from nle._pynethack.nethack import GEM_CLASS
+from nle._pynethack.nethack import GENOCIDED
+from nle._pynethack.nethack import GLYPH_BODY_OFF
+from nle._pynethack.nethack import GLYPH_CMAP_OFF
+from nle._pynethack.nethack import GLYPH_DETECT_OFF
+from nle._pynethack.nethack import GLYPH_EXPLODE_OFF
+from nle._pynethack.nethack import GLYPH_INVIS_OFF
+from nle._pynethack.nethack import GLYPH_INVISIBLE
+from nle._pynethack.nethack import GLYPH_MON_OFF
+from nle._pynethack.nethack import GLYPH_OBJ_OFF
+from nle._pynethack.nethack import GLYPH_PET_OFF
+from nle._pynethack.nethack import GLYPH_RIDDEN_OFF
+from nle._pynethack.nethack import GLYPH_STATUE_OFF
+from nle._pynethack.nethack import GLYPH_SWALLOW_OFF
+from nle._pynethack.nethack import GLYPH_WARNING_OFF
+from nle._pynethack.nethack import GLYPH_ZAP_OFF
+from nle._pynethack.nethack import ILLOBJ_CLASS
+from nle._pynethack.nethack import MAX_GLYPH
+from nle._pynethack.nethack import MAXEXPCHARS
+from nle._pynethack.nethack import MAXMCLASSES
+from nle._pynethack.nethack import MAXOCLASSES
+from nle._pynethack.nethack import MAXPCHARS
+from nle._pynethack.nethack import MAXWIN
+from nle._pynethack.nethack import MG_BW_LAVA
+from nle._pynethack.nethack import MG_CORPSE
+from nle._pynethack.nethack import MG_DETECT
+from nle._pynethack.nethack import MG_INVIS
+from nle._pynethack.nethack import MG_OBJPILE
+from nle._pynethack.nethack import MG_PET
+from nle._pynethack.nethack import MG_RIDDEN
+from nle._pynethack.nethack import MG_STATUE
+from nle._pynethack.nethack import NHW_MAP
+from nle._pynethack.nethack import NHW_MENU
+from nle._pynethack.nethack import NHW_MESSAGE
+from nle._pynethack.nethack import NHW_STATUS
+from nle._pynethack.nethack import NHW_TEXT
+from nle._pynethack.nethack import NLE_BL_AC
+from nle._pynethack.nethack import NLE_BL_ALIGN
+from nle._pynethack.nethack import NLE_BL_CAP
+from nle._pynethack.nethack import NLE_BL_CHA
+from nle._pynethack.nethack import NLE_BL_CON
+from nle._pynethack.nethack import NLE_BL_CONDITION
+from nle._pynethack.nethack import NLE_BL_DEPTH
+from nle._pynethack.nethack import NLE_BL_DEX
+from nle._pynethack.nethack import NLE_BL_DLEVEL
+from nle._pynethack.nethack import NLE_BL_DNUM
+from nle._pynethack.nethack import NLE_BL_ENE
+from nle._pynethack.nethack import NLE_BL_ENEMAX
+from nle._pynethack.nethack import NLE_BL_EXP
+from nle._pynethack.nethack import NLE_BL_GOLD
+from nle._pynethack.nethack import NLE_BL_HD
+from nle._pynethack.nethack import NLE_BL_HP
+from nle._pynethack.nethack import NLE_BL_HPMAX
+from nle._pynethack.nethack import NLE_BL_HUNGER
+from nle._pynethack.nethack import NLE_BL_INT
+from nle._pynethack.nethack import NLE_BL_SCORE
+from nle._pynethack.nethack import NLE_BL_STR25
+from nle._pynethack.nethack import NLE_BL_STR125
+from nle._pynethack.nethack import NLE_BL_TIME
+from nle._pynethack.nethack import NLE_BL_WIS
+from nle._pynethack.nethack import NLE_BL_X
+from nle._pynethack.nethack import NLE_BL_XP
+from nle._pynethack.nethack import NLE_BL_Y
+from nle._pynethack.nethack import NLE_BLSTATS_SIZE
+from nle._pynethack.nethack import NLE_INTERNAL_SIZE
+from nle._pynethack.nethack import NLE_INVENTORY_SIZE
+from nle._pynethack.nethack import NLE_INVENTORY_STR_LENGTH
+from nle._pynethack.nethack import NLE_MESSAGE_SIZE
+from nle._pynethack.nethack import NLE_MISC_SIZE
+from nle._pynethack.nethack import NLE_PROGRAM_STATE_SIZE
+from nle._pynethack.nethack import NLE_SCREEN_DESCRIPTION_LENGTH
+from nle._pynethack.nethack import NLE_TERM_CO
+from nle._pynethack.nethack import NLE_TERM_LI
+from nle._pynethack.nethack import NO_GLYPH
+from nle._pynethack.nethack import NUM_OBJECTS
+from nle._pynethack.nethack import NUM_ZAP
+from nle._pynethack.nethack import NUMMONS
+from nle._pynethack.nethack import OBJ_DESCR
+from nle._pynethack.nethack import OBJ_NAME
+from nle._pynethack.nethack import PANICKED
+from nle._pynethack.nethack import POISONING
+from nle._pynethack.nethack import POTION_CLASS
+from nle._pynethack.nethack import QUIT
+from nle._pynethack.nethack import RANDOM_CLASS
+from nle._pynethack.nethack import RING_CLASS
+from nle._pynethack.nethack import ROCK_CLASS
+from nle._pynethack.nethack import ROWNO
+from nle._pynethack.nethack import SCROLL_CLASS
+from nle._pynethack.nethack import SPBOOK_CLASS
+from nle._pynethack.nethack import STARVING
+from nle._pynethack.nethack import STONING
+from nle._pynethack.nethack import TILE_X
+from nle._pynethack.nethack import TILE_Y
+from nle._pynethack.nethack import TILE_Z
+from nle._pynethack.nethack import TOOL_CLASS
+from nle._pynethack.nethack import TRICKED
+from nle._pynethack.nethack import TURNED_SLIME
+from nle._pynethack.nethack import VENOM_CLASS
+from nle._pynethack.nethack import WAND_CLASS
+from nle._pynethack.nethack import WARNCOUNT
+from nle._pynethack.nethack import WEAPON_CLASS
+from nle._pynethack.nethack import class_sym
+from nle._pynethack.nethack import game_end_types
+from nle._pynethack.nethack import glyph2tile
+from nle._pynethack.nethack import glyph_is_body
+from nle._pynethack.nethack import glyph_is_cmap
+from nle._pynethack.nethack import glyph_is_detected_monster
+from nle._pynethack.nethack import glyph_is_invisible
+from nle._pynethack.nethack import glyph_is_monster
+from nle._pynethack.nethack import glyph_is_normal_monster
+from nle._pynethack.nethack import glyph_is_normal_object
+from nle._pynethack.nethack import glyph_is_object
+from nle._pynethack.nethack import glyph_is_pet
+from nle._pynethack.nethack import glyph_is_ridden_monster
+from nle._pynethack.nethack import glyph_is_statue
+from nle._pynethack.nethack import glyph_is_swallow
+from nle._pynethack.nethack import glyph_is_trap
+from nle._pynethack.nethack import glyph_is_warning
+from nle._pynethack.nethack import glyph_to_cmap
+from nle._pynethack.nethack import glyph_to_mon
+from nle._pynethack.nethack import glyph_to_obj
+from nle._pynethack.nethack import glyph_to_swallow
+from nle._pynethack.nethack import glyph_to_trap
+from nle._pynethack.nethack import glyph_to_warning
+from nle._pynethack.nethack import objclass
+from nle._pynethack.nethack import objdescr
+from nle._pynethack.nethack import permonst
+from nle._pynethack.nethack import symdef
 from nle.nethack.actions import *  # noqa: F403
 from nle.nethack.nethack import BLSTATS_SHAPE
 from nle.nethack.nethack import DUNGEON_SHAPE

nle/tests/test_dataset.py

@@ -4,16 +4,16 @@
 
 import numpy as np
 import pytest
-from test_converter import COLSROWS
-from test_converter import FINALFRAME
-from test_converter import FINALFRAMECOLORS
-from test_converter import TIMESTAMPS
-from test_converter import getfilename
-from test_db import conn  # noqa: F401
-from test_db import mockdata  # noqa: F401
 
 from nle.dataset import dataset
 from nle.dataset import db
+from nle.tests.test_converter import COLSROWS
+from nle.tests.test_converter import FINALFRAME
+from nle.tests.test_converter import FINALFRAMECOLORS
+from nle.tests.test_converter import TIMESTAMPS
+from nle.tests.test_converter import getfilename
+from nle.tests.test_db import conn  # noqa: F401
+from nle.tests.test_db import mockdata  # noqa: F401
 
 
 class TestDataset:

nle/tests/test_db.py

@@ -1,9 +1,10 @@
 import os
 
 import pytest
-import test_converter
 
 import nle.env.tasks
+import nle.nethack
+import nle.tests.test_converter as test_converter
 from nle.dataset import db
 from nle.dataset import populate_db
 

nle/tests/test_populate_db.py

@@ -1,11 +1,11 @@
 import json
 
 import pytest  # NOQA: F401
-from test_converter import getfilename
-from test_db import conn  # NOQA: F401
-from test_db import mockdata  # NOQA: F401
 
 from nle import nethack
+from nle.tests.test_converter import getfilename
+from nle.tests.test_db import conn  # NOQA: F401
+from nle.tests.test_db import mockdata  # NOQA: F401
 
 TTYRECS_TABLE_OFFSET = 0
 GAMES_TABLE_OFFSET = 5

nle/tests/test_tiles.py

@@ -2,7 +2,7 @@
 import numpy as np
 import pytest
 
-import nle
+import nle.nethack
 
 
 class TestTileset:

pyproject.toml

@@ -28,9 +28,16 @@ file = "LICENSE"
 [project.urls]
 Homepage = "https://github.com/NetHack-LE/nle"
 
+
 [build-system]
 # Lock build-deps as uv does not yet support locking of build deps: astral-sh/uv#5190
-requires = ["scikit-build-core~=0.10", "pybind11~=2.2", "setuptools-scm~=9.2.2"]
+requires = [
+  "scikit-build-core~=0.10",
+  "numpy",
+  "pybind11-stubgen",
+  "pybind11~=2.2",
+  "setuptools-scm~=9.2.2",
+]
 build-backend = "scikit_build_core.build"
 
 [tool.scikit-build]
@@ -42,6 +49,9 @@ generate = [
   { path = "nle/version.py", template = '__version__ = "${version}"' },
 ]
 
+#[tool.uv]
+#managed = false
+
 [tool.setuptools_scm]
 # This section is necessary to activate setuptools-scm, but can be empty
 
@@ -53,7 +63,7 @@ nle-ttyplay2 = "nle.scripts.ttyplay2:main"
 nle-read-tty = "nle.scripts.read_tty:main"
 
 [dependency-groups]
-dev = ["nle[dev]"]
+dev = ["nle[dev]", "pyrefly>=0.57.1"]
 
 [project.optional-dependencies]
 agent = ["torch>=1.3.1"]
@@ -115,3 +125,8 @@ before-all = "rm -rf {project}/build {project}/*.so {project}/CMakeCache.txt &&
 
 [tool.pytest.ini_options]
 testpaths = ["nle/tests"]
+
+[tool.pyrefly]
+disable-search-path-heuristics = true
+project-includes = ["."]
+project-excludes = ["src/**"]