Skip to content

Commit

Permalink
Fixed stub generation
Browse files Browse the repository at this point in the history
  • Loading branch information
themarpe committed May 28, 2022
1 parent e043880 commit 76e292f
Showing 1 changed file with 6 additions and 2 deletions.
8 changes: 6 additions & 2 deletions generate_stubs.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
MODULE_NAME = sys.argv[1]
DIRECTORY = sys.argv[2]

print(f'Generating stubs for module: {MODULE_NAME} in directory: {DIRECTORY}')
print(f'Generating stubs for module: "{MODULE_NAME}" in directory: "{DIRECTORY}"')

try:

Expand All @@ -33,14 +33,18 @@

# 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)

# Modify "*View" naming
nodes = re.findall('View\\[(\S*)\\]', final_stubs)
final_stubs = re.sub(r"View\[(\S*)\]", f'View_\\1', final_stubs)

# Writeout changes
file.seek(0)
file.write(final_stubs)
Expand Down

0 comments on commit 76e292f

Please sign in to comment.