Type more of the expansion functionality

Author: inducer

sumpy/expansion/__init__.py

@@ -28,7 +28,7 @@
 import logging
 from dataclasses import dataclass
 from itertools import accumulate
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, TypeAlias
 
 import numpy as np
 import sympy as sp
@@ -48,6 +48,13 @@
 logger = logging.getLogger(__name__)
 
 __doc__ = """
+.. autodata:: MultiIndex
+    :no-index:
+
+.. class:: MultiIndex
+
+    See above.
+
 Differential operator interface
 -------------------------------
 
@@ -58,14 +65,17 @@
 """
 
 
+MultiIndex: TypeAlias = tuple[int, ...]
+
+
 @dataclass(frozen=True)
 class DerivativeIdentifier:
     """
     .. autoattribute:: mi
     .. autoattribute: vec_idx
     """
 
-    mi: tuple[int, ...]
+    mi: MultiIndex
     """
     Multi-index of the derivative being taken, a tuple with a number of entries
     corresponding to the dimension.