target: fix ASTBuilder.get_function_declaration annotations

illinois-ceesd · Apr 13, 2024 · 52abd21 · 52abd21
1 parent d1437f3
commit 52abd21
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/loopy/target/__init__.py b/loopy/target/__init__.py
@@ -216,7 +216,7 @@ def get_function_definition(
     def get_function_declaration(
             self, codegen_state: CodeGenerationState,
             codegen_result: CodeGenerationResult, schedule_index: int
-            ) -> Tuple[Sequence[Tuple[str, str]], ASTType]:
+            ) -> Tuple[Sequence[Tuple[str, str]], Optional[ASTType]]:
         """Returns preambles and the AST for the function declaration."""
         raise NotImplementedError
 

diff --git a/loopy/target/pyopencl.py b/loopy/target/pyopencl.py
@@ -633,7 +633,7 @@ def generate_value_arg_setup(
 
     from genpy import If, Raise, Statement as S, Suite
 
-    result: List[str] = []
+    result: List[genpy.Generable] = []
     gen = result.append
 
     buf_indices_and_args = []
@@ -717,7 +717,7 @@ def generate_array_arg_setup(
     from loopy.kernel.array import ArrayBase
     from genpy import Statement as S, Suite
 
-    result: List[str] = []
+    result: List[genpy.Generable] = []
     gen = result.append
 
     cl_indices_and_args: List[Union[int, str]] = []
@@ -787,7 +787,7 @@ def get_function_definition(
     def get_function_declaration(
             self, codegen_state: CodeGenerationState,
             codegen_result: CodeGenerationResult, schedule_index: int
-            ) -> Tuple[Sequence[Tuple[str, str]], genpy.Generable]:
+            ) -> Tuple[Sequence[Tuple[str, str]], Optional[genpy.Generable]]:
         # no such thing in Python
         return [], None
 

diff --git a/loopy/target/python.py b/loopy/target/python.py
@@ -23,7 +23,7 @@
 THE SOFTWARE.
 """
 
-from typing import Tuple, Sequence
+from typing import Optional, Sequence, Tuple
 
 from pymbolic.mapper import Mapper
 from pymbolic.mapper.stringifier import StringifyMapper
@@ -169,7 +169,7 @@ def ast_module(self):
     def get_function_declaration(
             self, codegen_state: CodeGenerationState,
             codegen_result: CodeGenerationResult, schedule_index: int
-            ) -> Tuple[Sequence[Tuple[str, str]], None]:
+            ) -> Tuple[Sequence[Tuple[str, str]], Optional[Generable]]:
         return [], None
 
     def get_function_definition(self, codegen_state, codegen_result,