diff --git a/stellarphot/settings/models.py b/stellarphot/settings/models.py index 4dca97cd..0aa251b2 100644 --- a/stellarphot/settings/models.py +++ b/stellarphot/settings/models.py @@ -6,7 +6,15 @@ from .autowidgets import CustomBoundedIntTex -__all__ = ["ApertureSettings", "PhotometryFileSettings"] +from astropy.time import Time +from astropy.coordinates import SkyCoord +from astropy.units import Quantity +import astropy.units as u + +from ..core import QuantityType + + +__all__ = ["ApertureSettings", "PhotometryFileSettings", "Exoplanet"] class ApertureSettings(BaseModel): @@ -68,7 +76,8 @@ def outer_annulus(self): class PhotometryFileSettings(BaseModel): """ - An evolutionary step on the way to having a monolithic set of photometry settings. + An evolutionary step on the way to + having a monolithic set of photometry settings. """ image_folder: Path = Field( @@ -81,4 +90,55 @@ class PhotometryFileSettings(BaseModel): filter_pattern=["*.ecsv", "*.csv"], default="" ) -'fdsfafdafafdasdfasdfda' + +class TimeType(Time): + @classmethod + def __get_validators__(cls): + yield cls.validate + + @classmethod + def validate(cls, v): + return Time(v) + + +class SkyCoordType(SkyCoord): + @classmethod + def __get_validators__(cls): + yield cls.validate + + @classmethod + def validate(cls, v): + return SkyCoord(v) + + +class Exoplanet(BaseModel): + epoch: TimeType | None = None + period: QuantityType | None = None + identifier: str + coordinate: SkyCoordType + depth: float | None = None + duration: QuantityType | None = None + + class Config: + validate_all = True + validate_assignment = True + extra = "forbid" + json_encoders = { + Quantity: lambda v: f"{v.value} {v.unit}", + QuantityType: lambda v: f"{v.value} {v.unit}", + Time: lambda v: f"{v.value}", + } + + @classmethod + def validate_period(cls, values): + if u.get_physical_type(values["period"]) != "time": + raise ValueError( + f"Period does not have time units,currently has {values['period'].unit} units." + ) + + @classmethod + def validate_duration(cls, values): + if u.get_physical_type(values["duration"]) != "time": + raise ValueError( + f"Period does not have time units,currently has {values['duration'].unit} units." + ) diff --git a/stellarphot/settings/tests/test_models.py b/stellarphot/settings/tests/test_models.py index 1bef5e81..ff44024f 100644 --- a/stellarphot/settings/tests/test_models.py +++ b/stellarphot/settings/tests/test_models.py @@ -1,8 +1,12 @@ +from astropy.time import Time from pydantic import ValidationError import pytest from stellarphot.settings.models import ApertureSettings +from stellarphot.settings.models import Exoplanet +from astropy.coordinates import SkyCoord +import astropy.units as u DEFAULT_APERTURE_SETTINGS = dict(radius=5, gap=10, annulus_width=15) @@ -29,3 +33,25 @@ def test_create_invalid_values(bad_one): bad_settings[bad_one] = -1 with pytest.raises(ValidationError, match=bad_one): ApertureSettings(**bad_settings) + + +DEFAULT_EXOPLANET_SETTINGS = dict( + epoch=Time(0, format="jd"), + period=0 * u.min, + identifier="", + coordinate=SkyCoord( + ra="00:00:00.00", dec="+00:00:00.0", frame="icrs", unit=("hour", "degree") + ), + depth=0, + duration=0 * u.min, +) + + +def test_create_exoplanet_correctly(): + planet = Exoplanet(**DEFAULT_EXOPLANET_SETTINGS) + assert planet.epoch == DEFAULT_EXOPLANET_SETTINGS["epoch"] + assert planet.period == DEFAULT_EXOPLANET_SETTINGS["period"] + assert planet.identifier == DEFAULT_EXOPLANET_SETTINGS["identifier"] + assert planet.coordinate == DEFAULT_EXOPLANET_SETTINGS["coordinate"] + assert planet.depth == DEFAULT_EXOPLANET_SETTINGS["depth"] + assert planet.duration == DEFAULT_EXOPLANET_SETTINGS["duration"]