syntaxwebsite/app/util/membership.py

from app.models.user_membership import UserMembership
from app.enums.MembershipType import MembershipType
from app.extensions import db
from app.models.user import User
from datetime import datetime, timedelta

class UserHasHigherMembershipException(Exception):
    """
        Raised when a user has a higher membership than the one they are trying to be given.
    """
    pass
class UserDoesNotExistException(Exception):
    """
        Raised when a user does not exist.
    """
    pass
class UserDoesNotHaveMembershipException(Exception):
    """
        Raised when a user does not have a membership.
    """
    pass

def GetUserFromId( TargetUser : User | int ) -> User | None:
    """
        Gets a user from their ID.

        :param TargetUser: The user to get.

        :returns: User | None (The user, None if the user does not exist)
    """
    if type(TargetUser) is int:
        TargetUser = User.query.filter_by(id=TargetUser).first()
    return TargetUser

def GetUserMembership( TargetUser : User | int, changeToString : bool = False) -> MembershipType | None | str:
    """
        Gets the membership of a user.

        :param TargetUser: The user to get the membership of.
        :param changeToString: Whether to change the membership enum to a string or not.

        :returns: MembershipType | None | str (The membership of the user, None if the user does not exist, "None" if the user has no membership, or a string if changeToString is True)
    """
    TargetUser = GetUserFromId(TargetUser)
    if TargetUser is None:
        return None
    MembershipObj : UserMembership = UserMembership.query.filter_by(user_id=TargetUser.id).first()
    if MembershipObj is None:
        MembershipObj : UserMembership = UserMembership(
            TargetUser.id,
            MembershipType.NonBuildersClub,
            None,
            None
        )
        db.session.add(MembershipObj)
        db.session.commit()
    if MembershipObj.expiration is not None and MembershipObj.expiration < datetime.utcnow():
        MembershipObj.membership_type = MembershipType.NonBuildersClub
        MembershipObj.expiration = None
        MembershipObj.next_stipend = None
        MembershipObj.created = None
        db.session.commit()
    if changeToString:
        if MembershipObj.membership_type == MembershipType.NonBuildersClub:
            return "None"
        return MembershipObj.membership_type.name
    return MembershipObj.membership_type


def GiveUserMembership( TargetUser : User | int, Membership : MembershipType, expiration : timedelta = timedelta(days=31),ForceMembership : bool = False, throwException : bool = True, incrementExpirationIfSame : bool = True) -> None:
    """
        Gives a user a membership.

        :param TargetUser: The user to give the membership to.
        :param Membership: The membership to give the user.
        :param expiration: The expiration of the membership.
        :param ForceMembership: Whether to force the membership or not.
        :param throwException: Whether to throw an exception or not.
        :param incrementExpirationIfSame: Whether to increment the expiration if the user already has the membership.

        :returns: None
    """
    
    TargetUser = GetUserFromId(TargetUser)
    if TargetUser is None:
        return None
    CurrentMembership : MembershipType = GetUserMembership(TargetUser)
    if CurrentMembership is None:
        if throwException:
            raise UserDoesNotExistException("User does not exist.")
        return
    if CurrentMembership == Membership:
        if incrementExpirationIfSame:
            MembershipObj : UserMembership = UserMembership.query.filter_by(user_id=TargetUser.id).first()
            if MembershipObj is None:
                raise Exception("Membership object does not exist, this should never happen.")
            if MembershipObj.expiration is None:
                return
            MembershipObj.expiration = MembershipObj.expiration + expiration
            db.session.commit()
        return
    if CurrentMembership.value > Membership.value and not ForceMembership:
        if throwException:
            raise UserHasHigherMembershipException("User has a higher membership than the one they are trying to be given.")
        return
    MembershipObj : UserMembership = UserMembership.query.filter_by(user_id=TargetUser.id).first()
    if MembershipObj is None:
        raise Exception("Membership object does not exist, this should never happen.")
    MembershipObj.membership_type = Membership
    MembershipObj.created = datetime.utcnow()
    MembershipObj.expiration = datetime.utcnow() + expiration
    MembershipObj.next_stipend = datetime.utcnow()
    db.session.commit()

    return

def RemoveUserMembership( TargetUser : User | int ) -> None:
    """
        Removes a user's membership.

        :param TargetUser: The user to remove the membership of.

        :returns: None
    """

    TargetUser = GetUserFromId(TargetUser)
    if TargetUser is None:
        return None
    CurrentMembership : MembershipType = GetUserMembership(TargetUser)
    if CurrentMembership is MembershipType.NonBuildersClub:
        return
    MembershipObj : UserMembership = UserMembership.query.filter_by(user_id=TargetUser.id).first()
    if MembershipObj is None:
        raise Exception("Membership object does not exist, this should never happen.")
    MembershipObj.membership_type = MembershipType.NonBuildersClub
    MembershipObj.created = None
    MembershipObj.expiration = None
    MembershipObj.next_stipend = None
    db.session.commit()

    return
def IncrementExpirationLength( TargetUser : User | int, Length : timedelta = timedelta(days=31) ) -> None:
    """
        Increments the expiration length of a user's membership.

        :param TargetUser: The user to increment the expiration length of.
        :param Length: The length to increment the expiration by.

        :returns: None
    """
    TargetUser : User = GetUserFromId(TargetUser)
    if TargetUser is None:
        return None
    CurrentMembership : MembershipType = GetUserMembership(TargetUser)
    if CurrentMembership is None:
        raise UserDoesNotExistException("User does not exist.")
    if CurrentMembership == MembershipType.NonBuildersClub:
        raise UserDoesNotHaveMembershipException("User does not have a membership.")
    MembershipObj : UserMembership = UserMembership.query.filter_by(user_id=TargetUser.id).first()
    if MembershipObj is None:
        raise Exception("Membership object does not exist, this should never happen.")
    if MembershipObj.expiration is None:
        raise Exception("Membership object does not have an expiration, this should never happen.")
    MembershipObj.expiration = MembershipObj.expiration + Length
    db.session.commit()

    return