From f55908bc48f161d96bb74e407985173802c64cb7 Mon Sep 17 00:00:00 2001
From: Diogo <dgcruz983@gmail.com>
Date: Thu, 14 Dec 2023 15:58:45 -0500
Subject: [PATCH] Added stubs for python files generated from C++ (#136)

* added pybind11-stubgen

* docs for generating stubs

* added line to readme
---
 .gitignore           |  1 +
 README.md            |  7 +++++++
 docs/src/install.rst |  7 +++++++
 setup.py             | 34 +++++++++++++++++++++++++++++++---
 4 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/.gitignore b/.gitignore
index c382646a2..3a30aae9e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -11,6 +11,7 @@ __pycache__/
 venv/
 
 # Distribution / packaging
+python/mlx/core
 python/mlx/share
 python/mlx/include
 .Python
diff --git a/README.md b/README.md
index 1fe603552..057e95886 100644
--- a/README.md
+++ b/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
+```
\ No newline at end of file
diff --git a/docs/src/install.rst b/docs/src/install.rst
index 4ce1074a8..5244fac2f 100644
--- a/docs/src/install.rst
+++ b/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
 ^^^^^^^
 
diff --git a/setup.py b/setup.py
index ca833615c..4711a87c9 100644
--- a/setup.py
+++ b/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",
     )