Skip to content

Commit 6e910cb

Browse files
committed
Unified stubs generation and checking
1 parent 237e603 commit 6e910cb

File tree

4 files changed

+73
-28
lines changed

4 files changed

+73
-28
lines changed

CMakeLists.txt

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,22 @@ pybind11_add_module(${TARGET_NAME}
9797
src/log/LogBindings.cpp
9898
)
9999

100+
# Add stubs (pyi) generation step after building bindings
101+
execute_process(COMMAND "${PYTHON_EXECUTABLE}" "-c" "from mypy import api" RESULT_VARIABLE error OUTPUT_QUIET ERROR_QUIET)
102+
if(error)
103+
message(WARNING "Mypy not available - stubs won't be generated or checked")
104+
else()
105+
get_target_property(bindings_directory ${TARGET_NAME} LIBRARY_OUTPUT_DIRECTORY)
106+
if(NOT bindings_directory)
107+
set(bindings_directory ${CMAKE_CURRENT_BINARY_DIR})
108+
endif()
109+
message(STATUS "Mypy available, creating and checking stubs. Running with generate_stubs.py ${TARGET_NAME} ${bindings_directory}")
110+
add_custom_command(TARGET ${TARGET_NAME} POST_BUILD COMMAND
111+
${PYTHON_EXECUTABLE} "${CMAKE_CURRENT_LIST_DIR}/generate_stubs.py" "${TARGET_NAME}" "${bindings_directory}"
112+
DEPENDS "${CMAKE_CURRENT_LIST_DIR}/generate_stubs.py"
113+
WORKING_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}"
114+
)
115+
endif()
100116

101117
# Docstring options
102118
if(DEPTHAI_PYTHON_DOCSTRINGS_INPUT AND DEPTHAI_PYTHON_DOCSTRINGS_OUTPUT)

depthai-core

generate_stubs.py

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
import sys
2+
import subprocess
3+
import re
4+
import tempfile
5+
6+
# Usage
7+
if len(sys.argv) < 3:
8+
print(f"Usage: {sys.argv[0]} [module_name] [library_dir]")
9+
exit(-1)
10+
11+
MODULE_NAME = sys.argv[1]
12+
DIRECTORY = sys.argv[2]
13+
14+
print(f'Generating stubs for module: {MODULE_NAME} in directory: {DIRECTORY}')
15+
16+
try:
17+
18+
# Create stubs, add PYTHONPATH to find the build module
19+
# CWD to to extdir where the built module can be found to extract the types
20+
subprocess.check_call(['stubgen', '-p', MODULE_NAME, '-o', f'{DIRECTORY}'], cwd=DIRECTORY)
21+
22+
# Add py.typed
23+
open(f'{DIRECTORY}/depthai/py.typed', 'a').close()
24+
25+
# imports and overloads
26+
with open(f'{DIRECTORY}/depthai/__init__.pyi' ,'r+') as file:
27+
# Read
28+
contents = file.read()
29+
30+
# Add imports
31+
stubs_import = 'import depthai.node as node\nimport typing\nimport json\n' + contents
32+
# Create 'create' overloads
33+
nodes = re.findall('def \S*\(self\) -> node.(\S*):', stubs_import)
34+
overloads = ''
35+
for node in nodes:
36+
overloads = overloads + f'\\1@overload\\1def create(self, arg0: typing.Type[node.{node}]) -> node.{node}: ...'
37+
#print(f'{overloads}')
38+
final_stubs = re.sub(r"([\s]*)def create\(self, arg0: object\) -> Node: ...", f'{overloads}', stubs_import)
39+
40+
# Writeout changes
41+
file.seek(0)
42+
file.write(final_stubs)
43+
44+
# Flush previous stdout
45+
sys.stdout.flush()
46+
47+
# Check syntax
48+
with tempfile.NamedTemporaryFile() as config:
49+
config.write('[mypy]\nignore_errors = True\n'.encode())
50+
config.flush()
51+
subprocess.check_call([sys.executable, '-m' 'mypy', f'{DIRECTORY}/{MODULE_NAME}', f'--config-file={config.name}'])
52+
53+
except subprocess.CalledProcessError as err:
54+
exit(err.returncode)
55+
56+
exit(0)

setup.py

Lines changed: 0 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -173,33 +173,6 @@ def build_extension(self, ext):
173173
subprocess.check_call(['cmake', ext.sourcedir] + cmake_args, cwd=self.build_temp, env=env)
174174
subprocess.check_call(['cmake', '--build', '.'] + build_args, cwd=self.build_temp)
175175

176-
# Create stubs, add PYTHONPATH to find the build module
177-
# CWD to to extdir where the built module can be found to extract the types
178-
subprocess.check_call(['stubgen', '-p', MODULE_NAME, '-o', f'{extdir}'], cwd=extdir)
179-
180-
# Add py.typed
181-
open(f'{extdir}/depthai/py.typed', 'a').close()
182-
183-
# imports and overloads
184-
with open(f'{extdir}/depthai/__init__.pyi' ,'r+') as file:
185-
# Read
186-
contents = file.read()
187-
188-
# Add imports
189-
stubs_import = 'import depthai.node as node\nimport typing\nimport json\n' + contents
190-
# Create 'create' overloads
191-
nodes = re.findall('def \S*\(self\) -> node.(\S*):', stubs_import)
192-
overloads = ''
193-
for node in nodes:
194-
overloads = overloads + f'\\1@overload\\1def create(self, arg0: typing.Type[node.{node}]) -> node.{node}: ...'
195-
print(f'{overloads}')
196-
final_stubs = re.sub(r"([\s]*)def create\(self, arg0: object\) -> Node: ...", f'{overloads}', stubs_import)
197-
198-
# Writeout changes
199-
file.seek(0)
200-
file.write(final_stubs)
201-
202-
203176
setup(
204177
name=MODULE_NAME,
205178
version=__version__,

0 commit comments

Comments
 (0)