fix typing in make_loopy_program

Author: majosm

arraycontext/loopy.py

@@ -1,6 +1,33 @@
 """
 .. currentmodule:: arraycontext
 .. autofunction:: make_loopy_program
+
+References
+----------
+
+.. class:: InstructionBase
+
+    See :class:`loopy.InstructionBase`.
+
+.. class:: SubstitutionRule
+
+    See :class:`loopy.SubstitutionRule`.
+
+.. class:: ValueArg
+
+    See :class:`loopy.ValueArg`.
+
+.. class:: ArrayArg
+
+    See :class:`loopy.ArrayArg`.
+
+.. class:: TemporaryVariable
+
+    See :class:`loopy.TemporaryVariable`.
+
+.. class:: EllipsisType
+
+    See :data:`types.EllipsisType`.
 """
 from __future__ import annotations
 
@@ -42,7 +69,19 @@
 
 
 if TYPE_CHECKING:
-    from collections.abc import Mapping
+    from collections.abc import Mapping, Sequence
+    from types import EllipsisType
+
+    import islpy as isl
+
+    from loopy.kernel.data import (
+        ArrayArg,
+        SubstitutionRule,
+        TemporaryVariable,
+        ValueArg,
+    )
+    from loopy.kernel.instruction import InstructionBase
+    from pytools.tag import ToTagSetConvertible
 
 
 # {{{ loopy
@@ -52,8 +91,14 @@
         return_dict=True)
 
 
-def make_loopy_program(domains, statements, kernel_data=None,
-        name="mm_actx_kernel", tags=None):
+def make_loopy_program(
+        domains: str | Sequence[str | isl.BasicSet],
+        statements: str | Sequence[InstructionBase | SubstitutionRule | str],
+        kernel_data: Sequence[
+                ValueArg | ArrayArg | TemporaryVariable | EllipsisType | str
+            ] | None = None,
+        name: str = "mm_actx_kernel",
+        tags: ToTagSetConvertible = None):
     """Return a :class:`loopy.LoopKernel` suitable for use with
     :meth:`ArrayContext.call_loopy`.
     """