From 57f7b5f933ddbe6565f110944caba90cea4587d4 Mon Sep 17 00:00:00 2001 From: Corentin Cadiou Date: Mon, 27 Nov 2023 16:49:34 +0100 Subject: [PATCH] More extensive typing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Clément Robert --- yt/data_objects/time_series.py | 24 +++++++++--------------- 1 file changed, 9 insertions(+), 15 deletions(-) diff --git a/yt/data_objects/time_series.py b/yt/data_objects/time_series.py index ff64ee5320..0d7d626d7b 100644 --- a/yt/data_objects/time_series.py +++ b/yt/data_objects/time_series.py @@ -443,11 +443,9 @@ def particle_trajectories( def _get_by_attribute( self, attribute: str, - value: Union[unyt_quantity, tuple[float, str]], - tolerance: Union[None, unyt_quantity, tuple[float, str]] = None, - side: Union[ - Literal["nearest"], Literal["smaller"], Literal["larger"] - ] = "nearest", + value: Union[unyt_quantity, tuple[float, Union[Unit, str]]], + tolerance: Union[None, unyt_quantity, tuple[float, Union[Unit, str]]] = None, + side: Literal["nearest", "smaller", "larger"] = "nearest", ) -> "Dataset": r""" Get a dataset at or near to a given value. @@ -549,12 +547,10 @@ def _get_by_attribute( def get_by_time( self, - time: Union[unyt_quantity, tuple], - tolerance: Union[None, unyt_quantity, tuple] = None, - side: Union[ - Literal["nearest"], Literal["smaller"], Literal["larger"] - ] = "nearest", - ): + time: Union[unyt_quantity, tuple[float, Union[Unit, str]]], + tolerance: Union[None, unyt_quantity, tuple[float, Union[Unit, str]]] = None, + side: Literal["nearest", "smaller", "larger"] = "nearest", + ) -> Dataset: """ Get a dataset at or near to a given time. @@ -585,10 +581,8 @@ def get_by_redshift( self, redshift: float, tolerance: Optional[float] = None, - side: Union[ - Literal["nearest"], Literal["smaller"], Literal["larger"] - ] = "nearest", - ): + side: Literal["nearest", "smaller", "larger"] = "nearest", + ) -> Dataset: """ Get a dataset at or near to a given time.