Readd type hints

Author: daquinteroflex

tidy3d/plugins/smatrix/component_modelers/base.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import os
-import typing
 from abc import ABC, abstractmethod
 from typing import Optional, Union, get_args
 
@@ -23,6 +22,7 @@
 from tidy3d.plugins.smatrix.ports.modal import Port
 from tidy3d.plugins.smatrix.ports.rectangular_lumped import LumpedPort
 from tidy3d.plugins.smatrix.ports.wave import WavePort
+from tidy3d.web import Batch, BatchData
 
 # fwidth of gaussian pulse in units of central frequency
 FWIDTH_FRAC = 1.0 / 10
@@ -105,7 +105,7 @@ class AbstractComponentModeler(ABC, Tidy3dBaseModel):
     )
 
     # TODO properly refactor, plugins should not have web methods.
-    batch_cached: typing.Any = pd.Field(
+    batch_cached: BatchData = pd.Field(
         None,
         title="Batch (Cached)",
         description="DEPRECATED: Optional field to specify ``batch``. Only used as a workaround internally "
@@ -182,7 +182,7 @@ def to_file(self, fname: str) -> None:
         super(AbstractComponentModeler, self).to_file(fname=fname)  # noqa: UP008
 
     @cached_property
-    def batch(self):
+    def batch(self) -> Batch:
         """:class:`.Batch` associated with this component modeler."""
         # TODO properly refactor, plugins data types should not have web methods.
         from tidy3d.web.api.container import Batch
@@ -210,7 +210,7 @@ def batch_path(self) -> str:
         return self.batch._batch_path(path_dir=self.path_dir)
 
     @cached_property
-    def batch_data(self):
+    def batch_data(self) -> BatchData:
         """The :class:`.BatchData` associated with the simulations run for this component modeler."""
         return self.batch.run(path_dir=self.path_dir)
 
@@ -232,7 +232,7 @@ def _batch_path(self) -> str:
         """Where we store the batch for this :class:`AbstractComponentModeler` instance after the run."""
         return os.path.join(self.path_dir, "batch" + str(hash(self)) + ".hdf5")
 
-    def _run_sims(self, path_dir: str = DEFAULT_DATA_DIR):
+    def _run_sims(self, path_dir: str = DEFAULT_DATA_DIR) -> BatchData:
         """Run :class:`.Simulation` for each port and return the batch after saving."""
         _ = self.get_path_dir(path_dir)
         self.batch.to_file(self._batch_path)
@@ -247,11 +247,11 @@ def get_port_by_name(self, port_name: str) -> Port:
         return ports[0]
 
     @abstractmethod
-    def _construct_smatrix(self, batch_data) -> DataArray:
+    def _construct_smatrix(self, batch_data: BatchData) -> DataArray:
         """Post process :class:`.BatchData` to generate scattering matrix."""
 
     @abstractmethod
-    def _internal_construct_smatrix(self, batch_data) -> DataArray:
+    def _internal_construct_smatrix(self, batch_data: BatchData) -> DataArray:
         """Post process :class:`.BatchData` to generate scattering matrix, for internal use only."""
 
     def run(self, path_dir: str = DEFAULT_DATA_DIR) -> DataArray:
@@ -324,3 +324,6 @@ def sim_data_by_task_name(self, task_name: str) -> SimulationData:
         sim_data = self.batch_data[task_name]
         config.logging_level = log_level_cache
         return sim_data
+
+
+AbstractComponentModeler.update_forward_refs()

tidy3d/plugins/smatrix/component_modelers/modal.py

@@ -19,6 +19,7 @@
 from tidy3d.components.viz import add_ax_if_none, equal_aspect
 from tidy3d.exceptions import SetupError
 from tidy3d.plugins.smatrix.ports.modal import ModalPortDataArray, Port
+from tidy3d.web import BatchData
 
 from .base import FWIDTH_FRAC, AbstractComponentModeler
 
@@ -276,7 +277,7 @@ def _construct_smatrix(self) -> ModalPortDataArray:
         """Post process :class:`.BatchData` to generate scattering matrix."""
         return self._internal_construct_smatrix(batch_data=self.batch_data)
 
-    def _internal_construct_smatrix(self, batch_data) -> ModalPortDataArray:
+    def _internal_construct_smatrix(self, batch_data: BatchData) -> ModalPortDataArray:
         """Post process :class:`.BatchData` to generate scattering matrix, for internal use only."""
 
         max_mode_index_out, max_mode_index_in = self.max_mode_index

tidy3d/plugins/smatrix/component_modelers/terminal.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import os
-import typing
 from typing import Optional, Union
 
 import numpy as np
@@ -33,6 +32,7 @@
 from tidy3d.plugins.smatrix.ports.coaxial_lumped import CoaxialLumpedPort
 from tidy3d.plugins.smatrix.ports.rectangular_lumped import LumpedPort
 from tidy3d.plugins.smatrix.ports.wave import WavePort
+from tidy3d.web import BatchData
 
 
 class TerminalComponentModeler(AbstractComponentModeler):
@@ -317,7 +317,7 @@ def _construct_smatrix(self) -> TerminalPortDataArray:
 
         return construct_smatrix(self)
 
-    def _internal_construct_smatrix(self, batch_data: typing.Any = None) -> TerminalPortDataArray:
+    def _internal_construct_smatrix(self, batch_data: BatchData = None) -> TerminalPortDataArray:
         from tidy3d.plugins.smatrix.local_run import construct_smatrix
 
         return construct_smatrix(self)
@@ -529,3 +529,6 @@ def s_to_z(
         from tidy3d.plugins.smatrix.utils import s_to_z
 
         return s_to_z(s_matrix=s_matrix, reference=reference)
+
+
+TerminalComponentModeler.update_forward_refs()