From 6e910cb22dd00368cd62e3b4959319bcbea5d67f Mon Sep 17 00:00:00 2001
From: TheMarpe <martin.peterlin7@gmail.com>
Date: Fri, 10 Dec 2021 03:15:09 +0100
Subject: [PATCH] Unified stubs generation and checking

---
 CMakeLists.txt    | 16 ++++++++++++++
 depthai-core      |  2 +-
 generate_stubs.py | 56 +++++++++++++++++++++++++++++++++++++++++++++++
 setup.py          | 27 -----------------------
 4 files changed, 73 insertions(+), 28 deletions(-)
 create mode 100644 generate_stubs.py

diff --git a/CMakeLists.txt b/CMakeLists.txt
index a94ad9e06..b91cdde6f 100644
--- a/CMakeLists.txt
+++ b/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)
diff --git a/depthai-core b/depthai-core
index 5dda7d9ce..0b2e94516 160000
--- a/depthai-core
+++ b/depthai-core
@@ -1 +1 @@
-Subproject commit 5dda7d9ce52de3bee424689ded0689789246a023
+Subproject commit 0b2e94516920cc1e290ee22959d65342a8c2f6f9
diff --git a/generate_stubs.py b/generate_stubs.py
new file mode 100644
index 000000000..1ea1181df
--- /dev/null
+++ b/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)
diff --git a/setup.py b/setup.py
index 46035ad9f..52256c361 100644
--- a/setup.py
+++ b/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__,