Added stubs for python files generated from C++ (#136)

* added pybind11-stubgen

* docs for generating stubs

* added line to readme

.gitignore

@@ -11,6 +11,7 @@ __pycache__/
 venv/
 
 # Distribution / packaging
+python/mlx/core
 python/mlx/share
 python/mlx/include
 .Python

README.md

@@ -80,3 +80,10 @@ for more information on building the C++ and Python APIs from source.
 
 Check out the [contribution guidelines](CONTRIBUTING.md) for more information
 on contributing to MLX.
+
+Optional: Generate stubs for C++ python objects to enable IDE auto-completion
+
+```shell
+pip install ".[dev]"
+python setup.py generate_stubs
+```

docs/src/install.rst

@@ -88,6 +88,13 @@ To make sure the install is working run the tests with:
   pip install ".[testing]"
   python -m unittest discover python/tests
 
+Optional: Install stubs to enable auto completions and type checking from your IDE:
+
+.. code-block:: shell
+
+  pip install ".[dev]"
+  python setup.py generate_stubs
+
 C++ API
 ^^^^^^^
 

setup.py

@@ -9,7 +9,7 @@ import sysconfig
 from pathlib import Path
 from subprocess import run
 
-from setuptools import Extension, find_namespace_packages, setup
+from setuptools import Command, Extension, find_namespace_packages, setup
 from setuptools.command.build_ext import build_ext
 
 
@@ -125,6 +125,31 @@ class CMakeBuild(build_ext):
                     self.copy_tree(regular_dir, inplace_dir)
 
 
+class GenerateStubs(Command):
+    user_options = []
+
+    def initialize_options(self):
+        pass
+
+    def finalize_options(self):
+        pass
+
+    def run(self) -> None:
+        subprocess.run(["pybind11-stubgen", "mlx.core", "-o", "python"])
+        # Note, sed inplace on macos requires a backup prefix, delete the file after its generated
+        # this sed is needed to replace references from py::cpp_function to a generic Callable
+        subprocess.run(
+            [
+                "sed",
+                "-i",
+                "''",
+                "s/cpp_function/typing.Callable/g",
+                "python/mlx/core/__init__.pyi",
+            ]
+        )
+        subprocess.run(["rm", "python/mlx/core/__init__.pyi''"])
+
+
 # Read the content of README.md
 with open(Path(__file__).parent / "README.md", encoding="utf-8") as f:
     long_description = f.read()
@@ -150,9 +175,12 @@ if __name__ == "__main__":
         package_dir=package_dir,
         package_data=package_data,
         include_package_data=True,
-        extras_require={"testing": ["numpy", "torch"]},
+        extras_require={
+            "testing": ["numpy", "torch"],
+            "dev": ["pre-commit", "pybind11-stubgen"],
+        },
         ext_modules=[CMakeExtension("mlx.core")],
-        cmdclass={"build_ext": CMakeBuild},
+        cmdclass={"build_ext": CMakeBuild, "generate_stubs": GenerateStubs},
         zip_safe=False,
         python_requires=">=3.8",
     )