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

* added pybind11-stubgen

* docs for generating stubs

* added line to readme

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",
     )