Unified stubs generation and checking

CMakeLists.txt

@@ -97,6 +97,22 @@ pybind11_add_module(${TARGET_NAME}
     src/log/LogBindings.cpp
 )
 
+# Add stubs (pyi) generation step after building bindings
+execute_process(COMMAND "${PYTHON_EXECUTABLE}" "-c" "from mypy import api" RESULT_VARIABLE error OUTPUT_QUIET ERROR_QUIET)
+if(error)
+    message(WARNING "Mypy not available - stubs won't be generated or checked")
+else()
+    get_target_property(bindings_directory ${TARGET_NAME} LIBRARY_OUTPUT_DIRECTORY)
+    if(NOT bindings_directory)
+        set(bindings_directory ${CMAKE_CURRENT_BINARY_DIR})
+    endif()
+    message(STATUS "Mypy available, creating and checking stubs. Running with generate_stubs.py ${TARGET_NAME} ${bindings_directory}")
+    add_custom_command(TARGET ${TARGET_NAME} POST_BUILD COMMAND
+        ${PYTHON_EXECUTABLE} "${CMAKE_CURRENT_LIST_DIR}/generate_stubs.py" "${TARGET_NAME}" "${bindings_directory}"
+        DEPENDS "${CMAKE_CURRENT_LIST_DIR}/generate_stubs.py"
+        WORKING_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}"
+    )
+endif()
 
 # Docstring options
 if(DEPTHAI_PYTHON_DOCSTRINGS_INPUT AND DEPTHAI_PYTHON_DOCSTRINGS_OUTPUT)

generate_stubs.py

@@ -0,0 +1,56 @@
+import sys
+import subprocess
+import re
+import tempfile
+
+# Usage
+if len(sys.argv) < 3:
+    print(f"Usage: {sys.argv[0]} [module_name] [library_dir]")
+    exit(-1)
+
+MODULE_NAME = sys.argv[1]
+DIRECTORY = sys.argv[2]
+
+print(f'Generating stubs for module: {MODULE_NAME} in directory: {DIRECTORY}')
+
+try:
+
+    # 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'{DIRECTORY}'], cwd=DIRECTORY)
+
+    # Add py.typed
+    open(f'{DIRECTORY}/depthai/py.typed', 'a').close()
+
+    # imports and overloads
+    with open(f'{DIRECTORY}/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)
+
+    # Flush previous stdout
+    sys.stdout.flush()
+
+    # Check syntax
+    with tempfile.NamedTemporaryFile() as config:
+        config.write('[mypy]\nignore_errors = True\n'.encode())
+        config.flush()
+        subprocess.check_call([sys.executable, '-m' 'mypy', f'{DIRECTORY}/{MODULE_NAME}', f'--config-file={config.name}'])
+
+except subprocess.CalledProcessError as err:
+    exit(err.returncode)
+
+exit(0)

setup.py

@@ -173,33 +173,6 @@ 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=MODULE_NAME,
     version=__version__,