Add type annotations

inducer · Nov 14, 2023 · 1bb7bd3 · 1bb7bd3
1 parent debd306
commit 1bb7bd3
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/meshmode/discretization/poly_element.py b/meshmode/discretization/poly_element.py
@@ -28,6 +28,7 @@
 from warnings import warn
 
 import numpy as np
+from numpy.typing import ArrayLike
 from pytools import memoize_method, memoize_on_first_arg
 from meshmode.mesh import (
         MeshElementGroup as _MeshElementGroup,
@@ -41,6 +42,7 @@
         InterpolatoryElementGroupBase)
 
 import modepy as mp
+from modepy import Basis
 
 __doc__ = """
 Group types
@@ -518,8 +520,9 @@ def from_mesh_interp_matrix(self):
 
 class TensorProductElementGroupBase(PolynomialElementGroupBase,
         HypercubeElementGroupBase):
-    def __init__(self, mesh_el_group, order, index=None, *, basis,
-                 unit_nodes):
+    def __init__(self, mesh_el_group: _MeshTensorProductElementGroup,
+                 order: int, index=None, *, basis: Basis,
+                 unit_nodes: ArrayLike):
         """
         :arg basis: a :class:`modepy.TensorProductBasis`.
         :arg unit_nodes: unit nodes for the tensor product, obtained by