From b00263c817d79fbb2d500ec1b5faf48c81c515e1 Mon Sep 17 00:00:00 2001
From: Thomas Sileo <t@a4.io>
Date: Mon, 29 Jul 2019 22:46:53 +0200
Subject: [PATCH] More notification tweaks

---
 utils/meta.py          |  3 +++
 utils/notifications.py | 26 ++++++++++++++++++++------
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/utils/meta.py b/utils/meta.py
index 4b32283..b8c8c38 100644
--- a/utils/meta.py
+++ b/utils/meta.py
@@ -1,3 +1,4 @@
+from enum import unique
 from enum import Enum
 from typing import Any
 from typing import Dict
@@ -7,12 +8,14 @@ from little_boxes import activitypub as ap
 _SubQuery = Dict[str, Any]
 
 
+@unique
 class Box(Enum):
     INBOX = "inbox"
     OUTBOX = "outbox"
     REPLIES = "replies"
 
 
+@unique
 class MetaKey(Enum):
     NOTIFICATION = "notification"
     NOTIFICATION_UNREAD = "notification_unread"
diff --git a/utils/notifications.py b/utils/notifications.py
index 301e928..a984f34 100644
--- a/utils/notifications.py
+++ b/utils/notifications.py
@@ -26,7 +26,7 @@ def _is_from_outbox(activity: ap.BaseActivity) -> bool:
 
 def _flag_as_notification(activity: ap.BaseActivity, new_meta: _NewMeta) -> None:
     new_meta.update(
-        **{_meta(MetaKey.NOTIFICATION): True, _meta(MetaKey.NOTIFICATION_UNREAD): True}
+        {_meta(MetaKey.NOTIFICATION): True, _meta(MetaKey.NOTIFICATION_UNREAD): True}
     )
     return None
 
@@ -58,7 +58,7 @@ def _accept_set_inbox_flags(activity: ap.Accept, new_meta: _NewMeta) -> None:
 
     # This Accept will be a "You started following $actor" notification
     _flag_as_notification(activity, new_meta)
-    new_meta.update(**{_meta(MetaKey.NOTIFICATION_FOLLOWS_BACK): follows_back})
+    new_meta.update({_meta(MetaKey.NOTIFICATION_FOLLOWS_BACK): follows_back})
     return None
 
 
@@ -83,7 +83,7 @@ def _follow_set_inbox_flags(activity: ap.Follow, new_meta: _NewMeta) -> None:
 
     # This Follow will be a "$actor started following you" notification
     _flag_as_notification(activity, new_meta)
-    new_meta.update(**{_meta(MetaKey.NOTIFICATION_FOLLOWS_BACK): follows_back})
+    new_meta.update({_meta(MetaKey.NOTIFICATION_FOLLOWS_BACK): follows_back})
     return None
 
 
@@ -98,7 +98,7 @@ def _like_set_inbox_flags(activity: ap.Like, new_meta: _NewMeta) -> None:
         Tasks.cache_object(activity.id)
 
         # Also set the "keep mark" for the GC (as we want to keep it forever)
-        new_meta.update(**{_meta(MetaKey.GC_KEEP): True})
+        new_meta.update({_meta(MetaKey.GC_KEEP): True})
 
     return None
 
@@ -111,9 +111,23 @@ def _announce_set_inbox_flags(activity: ap.Announce, new_meta: _NewMeta) -> None
         _flag_as_notification(activity, new_meta)
 
         # Also set the "keep mark" for the GC (as we want to keep it forever)
-        new_meta.update(**{_meta(MetaKey.GC_KEEP): True})
+        new_meta.update({_meta(MetaKey.GC_KEEP): True})
 
-    # Cache the object in all case (for display on the notifcation page)
+    # Cache the object in all case (for display on the notifcation page **and** the stream page)
     Tasks.cache_object(activity.id)
 
     return None
+
+
+@set_inbox_flags.register
+def _undo_set_inbox_flags(activity: ap.Undo, new_meta: _NewMeta) -> None:
+    obj = activity.get_object()
+
+    if obj.has_type(ap.ActivityType.FOLLOW):
+        # Flag it as a noticiation (for the "$actor unfollowed you"
+        _flag_as_notification(activity, new_meta)
+
+        # Also set the "keep mark" for the GC (as we want to keep it forever)
+        new_meta.update({_meta(MetaKey.GC_KEEP): True})
+
+    return None