From ba5ef916910403de96e8bfc214d29a59e5c35b27 Mon Sep 17 00:00:00 2001 From: purplenicole730 Date: Fri, 24 Jan 2025 14:10:52 -0500 Subject: [PATCH 1/2] import resource base modules to grab additional functions --- cli/module_generate/scripts/generate_stubs.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/cli/module_generate/scripts/generate_stubs.py b/cli/module_generate/scripts/generate_stubs.py index 72a076ff55a..83946f83d4f 100644 --- a/cli/module_generate/scripts/generate_stubs.py +++ b/cli/module_generate/scripts/generate_stubs.py @@ -155,6 +155,24 @@ def main( ) abstract_methods.append(indented_code) + type_module = import_module(f"viam.{resource_type}s.{resource_type}_base") + with open(type_module.__file__, "r") as f: + tree = ast.parse(f.read()) + for stmt in tree.body: + if isinstance(stmt, ast.ClassDef): + for cstmt in stmt.body: + if isinstance(cstmt, ast.AsyncFunctionDef): + replace_async_func("", cstmt, []) + indented_code = '\n'.join( + [' ' + line for line in ast.unparse(cstmt).splitlines()] + ) + abstract_methods.append(indented_code) + if cstmt.name == "do_command": + imports.append("from viam.utils import ValueTypes") + elif cstmt.name == "get_geometries": + imports.append("from typing import List") + imports.append("from viam.proto.common import Geometry") + model_name_pascal = "".join( [word.capitalize() for word in slugify(model_name).split("-")] ) From ecd3ac7f8246f9373de107c91690a5aa4214734d Mon Sep 17 00:00:00 2001 From: purplenicole730 Date: Fri, 24 Jan 2025 15:37:26 -0500 Subject: [PATCH 2/2] add optional as import --- cli/module_generate/scripts/generate_stubs.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/cli/module_generate/scripts/generate_stubs.py b/cli/module_generate/scripts/generate_stubs.py index 83946f83d4f..3d3fbb15c86 100644 --- a/cli/module_generate/scripts/generate_stubs.py +++ b/cli/module_generate/scripts/generate_stubs.py @@ -3,7 +3,7 @@ import subprocess import sys from importlib import import_module -from typing import List, Set +from typing import List, Set, Union def return_attribute(value: str, attr: str) -> ast.Attribute: @@ -168,9 +168,10 @@ def main( ) abstract_methods.append(indented_code) if cstmt.name == "do_command": + imports.append("from typing import Optional") imports.append("from viam.utils import ValueTypes") elif cstmt.name == "get_geometries": - imports.append("from typing import List") + imports.append("from typing import List, Optional") imports.append("from viam.proto.common import Geometry") model_name_pascal = "".join(