diff --git a/bot/models.py b/bot/models.py index 1be06af..f9acc3a 100644 --- a/bot/models.py +++ b/bot/models.py @@ -1,5 +1,6 @@ import datetime import enum +import functools from sqlalchemy import Boolean, Column, DateTime, Enum, ForeignKey, Integer, Text from sqlalchemy.ext.declarative import declarative_base @@ -36,7 +37,7 @@ class Guild(Base): submission_channel = Column(Integer, nullable=True) # The channel being scanned for messages by this particular guild. current_period = Column(Integer, ForeignKey('period.id'), nullable=True) # The period currently active for this guild. - all_periods = relationship("Period", back_populates="guild") # All periods ever started inside this guild + periods = relationship("Period", back_populates="guild") # All periods ever started inside this guild active = Column(Boolean, default=True) # Whether or not the bot is active inside the given Guild. Used for better querying. joined = Column(DateTime, default=datetime.datetime.utcnow) # The initial join time for this bot to a particular Discord. @@ -73,15 +74,24 @@ class Period(Base): voting_time = Column(DateTime, nullable=True) # When this period switched to the Voting state. finished_time = Column(DateTime, nullable=True) # When this period switched to the Finished state. + def check_not_finished(self, func): + """ + Throws `FinishedPeriod` if the period has already completed, is inactive, or is in it's Finished State. + """ + @functools.wraps(func) + def wrapper(*args, **kwargs): + if self.state is PeriodStates.FINISHED: raise FinishedPeriod(f"Period is in it's Finished state.") + elif not self.active: raise FinishedPeriod("Period is no longer active.") + elif self.completed: raise FinishedPeriod("Period is already completed.") + func(*args, **kwargs) + + return wrapper + + @check_not_finished def advance_state(self) -> PeriodStates: """ Advances the current recorded state of this Period, recording timestamps as needed. - - Throws FinishedPeriod if the period has already completed. """ - if self.state is PeriodStates.FINISHED: raise FinishedPeriod(f"Period is in it's Finished state.") - elif not self.active: raise FinishedPeriod("Period is no longer active.") - next_state = PeriodStates(int(self.state) + 1) if self.state == PeriodStates.READY: self.submissions_time = datetime.datetime.utcnow() @@ -95,3 +105,13 @@ class Period(Base): self.active = False return next_state + + @check_not_finished + def deactivate(self) -> None: + """ + Deactivates the period, setting it as inactive. + + Use `advance_state` if you want to properly advance the state. + """ + self.finished_time = datetime.datetime.utcnow() + self.active = False