Skip to content

Commit

Permalink
fix some type error
Browse files Browse the repository at this point in the history
  • Loading branch information
@parkma99
parkma99 committed Mar 9, 2022
1 parent 51d9ecc commit 608af2a
Showing 1 changed file with 27 additions and 25 deletions.
52 changes: 27 additions & 25 deletions stingray/base.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,15 @@
from astropy.time import Time, TimeDelta
from astropy.units import Quantity

from typing import TYPE_CHECKING, Union
from typing import TYPE_CHECKING, Type, TypeVar, Union
if TYPE_CHECKING:
from xarray import Dataset
from pandas import DataFrame
from astropy.timeseries import TimeSeries
from astropy.time import TimeDelta
TTime = Union[Time, TimeDelta, Quantity, np.array]
import numpy.typing as npt
TTime = Union[Time, TimeDelta, Quantity, npt.ArrayLike]
Tso = TypeVar("Tso", bound = "StingrayObject")


class StingrayObject(object):
Expand Down Expand Up @@ -118,7 +120,7 @@ def to_astropy_table(self) -> Table:
return ts

@classmethod
def from_astropy_table(cls, ts: Table) -> StingrayObject:
def from_astropy_table(cls: Type[Tso], ts: Table) -> Tso:
"""Create a Stingray Object object from data in an Astropy Table.
The table MUST contain at least a column named like the
Expand All @@ -140,11 +142,11 @@ def from_astropy_table(cls, ts: Table) -> StingrayObject:
array_attrs = ts.colnames

# Set the main attribute first
mainarray = np.array(ts[cls.main_array_attr])
setattr(cls, cls.main_array_attr, mainarray)
mainarray = np.array(ts[cls.main_array_attr]) # type: ignore
setattr(cls, cls.main_array_attr, mainarray) # type: ignore

for attr in array_attrs:
if attr == cls.main_array_attr:
if attr == cls.main_array_attr: # type: ignore
continue
setattr(cls, attr.lower(), np.array(ts[attr]))

Expand Down Expand Up @@ -175,7 +177,7 @@ def to_xarray(self) -> Dataset:
return ts

@classmethod
def from_xarray(cls, ts: Dataset) -> StingrayObject:
def from_xarray(cls: Type[Tso], ts: Dataset) -> Tso:
"""Create a `StingrayObject` from data in an xarray Dataset.
The dataset MUST contain at least a column named like the
Expand All @@ -190,18 +192,18 @@ def from_xarray(cls, ts: Dataset) -> StingrayObject:
"""
cls = cls()

if len(ts[cls.main_array_attr]) == 0:
if len(ts[cls.main_array_attr]) == 0: # type: ignore
# return an empty object
return cls

array_attrs = ts.coords

# Set the main attribute first
mainarray = np.array(ts[cls.main_array_attr])
setattr(cls, cls.main_array_attr, mainarray)
mainarray = np.array(ts[cls.main_array_attr]) # type: ignore
setattr(cls, cls.main_array_attr, mainarray) # type: ignore

for attr in array_attrs:
if attr == cls.main_array_attr:
if attr == cls.main_array_attr: # type: ignore
continue
setattr(cls, attr, np.array(ts[attr]))

Expand Down Expand Up @@ -233,7 +235,7 @@ def to_pandas(self) -> DataFrame:
return ts

@classmethod
def from_pandas(cls, ts: DataFrame) -> StingrayObject:
def from_pandas(cls: Type[Tso], ts: DataFrame) -> Tso:
"""Create an `StingrayObject` object from data in a pandas DataFrame.
The dataframe MUST contain at least a column named like the
Expand All @@ -256,11 +258,11 @@ def from_pandas(cls, ts: DataFrame) -> StingrayObject:
array_attrs = ts.columns

# Set the main attribute first
mainarray = np.array(ts[cls.main_array_attr])
setattr(cls, cls.main_array_attr, mainarray)
mainarray = np.array(ts[cls.main_array_attr]) # type: ignore
setattr(cls, cls.main_array_attr, mainarray) # type: ignore

for attr in array_attrs:
if attr == cls.main_array_attr:
if attr == cls.main_array_attr: # type: ignore
continue
setattr(cls, attr, np.array(ts[attr]))

Expand All @@ -271,7 +273,7 @@ def from_pandas(cls, ts: DataFrame) -> StingrayObject:
return cls

@classmethod
def read(cls, filename: str, fmt: str = None, format_=None) -> StingrayObject:
def read(cls: Type[Tso], filename: str, fmt: str = None, format_=None) -> Tso:
r"""Generic reader for :class`StingrayObject`
Currently supported formats are
Expand Down Expand Up @@ -415,7 +417,7 @@ def write(self, filename: str, fmt: str = None, format_=None) -> None:


class StingrayTimeseries(StingrayObject):
def to_astropy_timeseries(self) -> StingrayTimeseries:
def to_astropy_timeseries(self) -> TimeSeries:
"""Save the ``StingrayTimeseries`` to an ``Astropy`` timeseries.
Array attributes (time, pi, energy, etc.) are converted
Expand Down Expand Up @@ -443,8 +445,8 @@ def to_astropy_timeseries(self) -> StingrayTimeseries:
if data == {}:
data = None

if self.time is not None and np.size(self.time) > 0:
times = TimeDelta(self.time * u.s)
if self.time is not None and np.size(self.time) > 0: # type: ignore
times = TimeDelta(self.time * u.s) # type: ignore
ts = TimeSeries(data=data, time=times)
else:
ts = TimeSeries()
Expand Down Expand Up @@ -483,7 +485,7 @@ def from_astropy_timeseries(cls, ts: TimeSeries) -> StingrayTimeseries:
mjdref = ts.meta["mjdref"]

time, mjdref = interpret_times(time, mjdref)
cls.time = np.asarray(time)
cls.time = np.asarray(time) # type: ignore

array_attrs = ts.colnames
for key, val in ts.meta.items():
Expand Down Expand Up @@ -512,10 +514,10 @@ def change_mjdref(self, new_mjdref: float) -> StingrayTimeseries:
new_lc : :class:`StingrayTimeseries` object
The new time series, shifted by MJDREF
"""
time_shift = (self.mjdref - new_mjdref) * 86400
time_shift = (self.mjdref - new_mjdref) * 86400 # type: ignore

ts = self.shift(time_shift)
ts.mjdref = new_mjdref
ts.mjdref = new_mjdref # type: ignore
return ts

def shift(self, time_shift: float) -> StingrayTimeseries:
Expand All @@ -534,14 +536,14 @@ def shift(self, time_shift: float) -> StingrayTimeseries:
"""
ts = copy.deepcopy(self)
ts.time = np.asarray(ts.time) + time_shift
ts.time = np.asarray(ts.time) + time_shift # type: ignore
if hasattr(ts, "gti"):
ts.gti = np.asarray(ts.gti) + time_shift
ts.gti = np.asarray(ts.gti) + time_shift # type: ignore

return ts


def interpret_times(time: TTime, mjdref: float = 0) -> tuple[np.array, float]:
def interpret_times(time: TTime, mjdref: float = 0) -> tuple[npt.ArrayLike, float]:
"""Understand the format of input times, and return seconds from MJDREF
Parameters
Expand Down

0 comments on commit 608af2a

Please sign in to comment.