Add type hints for get_eigenmode_coefficients and fix type hint in Medium

python/geom.py

@@ -10,7 +10,7 @@
 import math
 from numbers import Number
 import operator
-from typing import List, NamedTuple, Optional, Type, Tuple, Union
+from typing import List, NamedTuple, Optional, Tuple, Union
 import warnings
 
 import numpy as np
@@ -321,8 +321,8 @@ def __init__(
         epsilon_offdiag: Vector3 = Vector3(),
         mu_diag: Vector3 = Vector3(1.0, 1.0, 1.0),
         mu_offdiag: Vector3 = Vector3(),
-        E_susceptibilities: Optional[List[Type[Susceptibility]]] = None,
-        H_susceptibilities: Optional[List[Type[Susceptibility]]] = None,
+        E_susceptibilities: Optional[List[Susceptibility]] = None,
+        H_susceptibilities: Optional[List[Susceptibility]] = None,
         E_chi2_diag: Vector3 = Vector3(),
         E_chi3_diag: Vector3 = Vector3(),
         H_chi2_diag: Vector3 = Vector3(),

python/simulation.py

@@ -1,3 +1,7 @@
+"""
+A collection of objects and helper routines for setting up the simulation.
+"""
+
 import functools
 import math
 import numbers
@@ -8,7 +12,7 @@
 import sys
 import warnings
 from collections import OrderedDict, namedtuple
-from typing import Callable, List, Optional, Tuple, Union
+from typing import Callable, List, NamedTuple, Optional, Tuple, Union
 from scipy.interpolate import pade
 
 try:
@@ -4125,18 +4129,18 @@ def get_array_slice_dimensions(self, component, vol=None, center=None, size=None
 
     def get_eigenmode_coefficients(
         self,
-        flux,
-        bands,
-        eig_parity=mp.NO_PARITY,
-        eig_vol=None,
-        eig_resolution=0,
-        eig_tolerance=1e-12,
-        kpoint_func=None,
-        direction=mp.AUTOMATIC,
-    ):
+        flux: DftFlux,
+        bands: Union[List[int], DiffractedPlanewave],
+        eig_parity: int = mp.NO_PARITY,
+        eig_vol: Volume = None,
+        eig_resolution: float = 0,
+        eig_tolerance: float = 1e-12,
+        kpoint_func: Callable[[float, int], float] = None,
+        direction: int = mp.AUTOMATIC,
+    ) -> NamedTuple:
         """
-        Given a flux object and list of band indices `bands` or `DiffractedPlanewave`, return a `namedtuple` with the
-        following fields:
+        Given a flux object and list of band indices (integers) `bands` or a `DiffractedPlanewave` object,
+        return a `namedtuple` with the following fields:
 
         + `alpha`: the complex eigenmode coefficients as a 3d NumPy array of size
           (`len(bands)`, `flux.nfreqs`, `2`). The last/third dimension refers to modes