Added typing information

themarpe · themarpe · commit cfb227cdf53c · 2021-11-03T19:18:54.000+01:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,2 +1,2 @@
 [build-system]
-requires = ["setuptools", "wheel"] # Must be preinstalled "cmake>=3.2.0"
+requires = ["setuptools", "wheel", "mypy"] # Must be preinstalled "cmake>=3.2.0"
diff --git a/setup.py b/setup.py
@@ -11,6 +11,8 @@
 from setuptools.command.build_ext import build_ext
 from distutils.version import LooseVersion
 
+### NAME
+MODULE_NAME = 'depthai'
 
 ### VERSION
 here = os.path.abspath(os.path.dirname(__file__))
@@ -171,10 +173,35 @@ def build_extension(self, ext):
         subprocess.check_call(['cmake', ext.sourcedir] + cmake_args, cwd=self.build_temp, env=env)
         subprocess.check_call(['cmake', '--build', '.'] + build_args, cwd=self.build_temp)
 
+        # Create stubs, add PYTHONPATH to find the build module
+        # CWD to to extdir where the built module can be found to extract the types
+        subprocess.check_call(['stubgen', '-p', MODULE_NAME, '-o', f'{extdir}'], cwd=extdir)
+
+        # Add py.typed
+        open(f'{extdir}/depthai/py.typed', 'a').close()
+
+        # imports and overloads
+        with open(f'{extdir}/depthai/__init__.pyi' ,'r+') as file:
+            # Read
+            contents = file.read()
+
+            # Add imports
+            stubs_import = 'import depthai.node as node\nimport typing\nimport json\n' + contents
+            # Create 'create' overloads
+            nodes = re.findall('def \S*\(self\) -> node.(\S*):', stubs_import)
+            overloads = ''
+            for node in nodes:
+                overloads = overloads + f'\\1@overload\\1def create(self, arg0: typing.Type[node.{node}]) -> node.{node}: ...'
+            print(f'{overloads}')
+            final_stubs = re.sub(r"([\s]*)def create\(self, arg0: object\) -> Node: ...", f'{overloads}', stubs_import)
+
+            # Writeout changes
+            file.seek(0)
+            file.write(final_stubs)
 
 
 setup(
-    name='depthai',
+    name=MODULE_NAME,
     version=__version__,
     author='Luxonis',
     author_email='support@luxonis.com',
@@ -183,7 +210,7 @@ def build_extension(self, ext):
     long_description=long_description,
     long_description_content_type="text/markdown",
     url="https://github.com/luxonis/depthai-python",
-    ext_modules=[CMakeExtension('depthai')],
+    ext_modules=[CMakeExtension(MODULE_NAME)],
     cmdclass={
         'build_ext': CMakeBuild
     },

Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,2 @@`
`1`	`1`	`[build-system]`
`2`		`-requires = ["setuptools", "wheel"] # Must be preinstalled "cmake>=3.2.0"`
	`2`	`+requires = ["setuptools", "wheel", "mypy"] # Must be preinstalled "cmake>=3.2.0"`