<s> from modularodm import fields <EOL> from framework . mongo import StoredObject , ObjectId <EOL> from modularodm . exceptions import ValidationValueError <EOL> from website . project . model import Node <EOL> from website . notifications . constants import NOTIFICATION_TYPES <EOL> def validate_subscription_type ( value ) : <EOL> if value not in NOTIFICATION_TYPES : <EOL> raise ValidationValueError <EOL> class NotificationSubscription ( StoredObject ) : <EOL> _id = fields . StringField ( primary = True ) <EOL> event_name = fields . StringField ( ) <EOL> owner = fields . AbstractForeignField ( ) <EOL> none = fields . ForeignField ( '<STR_LIT:user>' , list = True ) <EOL> email_digest = fields . ForeignField ( '<STR_LIT:user>' , list = True ) <EOL> email_transactional = fields . ForeignField ( '<STR_LIT:user>' , list = True ) <EOL> def add_user_to_subscription ( self , user , notification_type , save = True ) : <EOL> for nt in NOTIFICATION_TYPES : <EOL> if user in getattr ( self , nt ) : <EOL> if nt != notification_type : <EOL> getattr ( self , nt ) . remove ( user ) <EOL> else : <EOL> if nt == notification_type : <EOL> getattr ( self , nt ) . append ( user ) <EOL> if notification_type != '<STR_LIT:none>' and isinstance ( self . owner , Node ) and self . owner . parent_node : <EOL> user_subs = self . owner . parent_node . child_node_subscriptions <EOL> if self . owner . _id not in user_subs . setdefault ( user . _id , [ ] ) : <EOL> user_subs [ user . _id ] . append ( self . owner . _id ) <EOL> self . owner . parent_node . save ( ) <EOL> if save : <EOL> self . save ( ) <EOL> def remove_user_from_subscription ( self , user , save = True ) : <EOL> for notification_type in NOTIFICATION_TYPES : <EOL> try : <EOL> getattr ( self , notification_type , [ ] ) . remove ( user ) <EOL> except ValueError : <EOL> pass <EOL> if isinstance ( self . owner , Node ) and self . owner . parent_node : <EOL> try : <EOL> self . owner . parent_node . child_node_subscriptions . get ( user . _id , [ ] ) . remove ( self . owner . _id ) <EOL> self . owner . parent_node . save ( ) <EOL> except ValueError : <EOL> pass <EOL> if save : <EOL> self . save ( ) <EOL> class NotificationDigest ( StoredObject ) : <EOL> _id = fields . StringField ( primary = True , default = lambda : str ( ObjectId ( ) ) ) <EOL> user_id = fields . StringField ( index = True ) <EOL> timestamp = fields . DateTimeField ( ) <EOL> send_type = fields . StringField ( index = True , validate = validate_subscription_type ) <EOL> event = fields . StringField ( ) <EOL> message = fields . StringField ( ) <EOL> node_lineage = fields . StringField ( list = True ) </s>
