From b14371ce9716ce527bf3a723bb0bb0aabe942e2e Mon Sep 17 00:00:00 2001
From: Mikael Dallaire Cote <110583667+0mdc@users.noreply.github.com>
Date: Tue, 25 Jun 2024 15:03:04 -0400
Subject: [PATCH] Add stubs to python bindings to provide typing information.

---
 requirements.txt                |  1 +
 setup.py                        | 21 +++++++++++++++++++--
 src_python/habitat_sim/py.typed |  0
 3 files changed, 20 insertions(+), 2 deletions(-)
 create mode 100644 src_python/habitat_sim/py.typed

diff --git a/requirements.txt b/requirements.txt
index 3985aaf148..2018519c57 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,5 +7,6 @@ numba
 numpy==1.26.4
 numpy-quaternion
 pillow
+pybind11-stubgen
 scipy>=1.10.1
 tqdm
diff --git a/setup.py b/setup.py
index f1434405ab..3c2d7119a7 100644
--- a/setup.py
+++ b/setup.py
@@ -264,6 +264,16 @@ def run(self):
         for ext in self.extensions:
             self.build_extension(ext)
 
+        # Build stubs for 'habitat_sim_bindings' to provide type information.
+        subprocess.run(
+            [
+                "pybind11-stubgen",
+                "habitat_sim._ext.habitat_sim_bindings",
+                "-o",
+                "src_python",
+            ]
+        )
+
     def build_extension(self, ext):
         extdir = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.name)))
 
@@ -450,21 +460,28 @@ def load(filename):
     setup(
         name="habitat_sim",
         version=habitat_sim.__version__,
-        author="FAIR A-STAR",
+        author="Meta",
         description="A high performance simulator for training embodied agents",
         long_description="",
         packages=find_packages(where="src_python"),
         package_dir={"": "src_python"},
         install_requires=requirements,
         tests_require=["hypothesis", "pytest-benchmark", "pytest"],
+        setup_requires=["pybind11-stubgen"],
         python_requires=">=3.9",
         # add extension module
         ext_modules=[CMakeExtension("habitat_sim._ext.habitat_sim_bindings", "src")],
         # add custom build_ext command
-        cmdclass=dict(build_ext=CMakeBuild),
+        cmdclass={
+            "build_ext": CMakeBuild,
+        },
         zip_safe=False,
         include_package_data=True,
+        package_data={
+            "habitat_sim": ["py.typed", "**/*.pyi"],
+        },
     )
+
     pymagnum_build_dir = osp.join(
         _cmake_build_dir, "deps", "magnum-bindings", "src", "python"
     )
diff --git a/src_python/habitat_sim/py.typed b/src_python/habitat_sim/py.typed
new file mode 100644
index 0000000000..e69de29bb2