move pages.count_followers to User, use it in activitypub.actor

activitypub.py

@@ -927,18 +927,13 @@ def follower_collection(id, collection):
         return page, {'Content-Type': as2.CONTENT_TYPE}
 
     # collection
-    prop = Follower.to if collection == 'followers' else Follower.from_
-    count = Follower.query(
-        Follower.status == 'active',
-        prop == g.user.key,
-    ).count()  # TODO: cache, unify with pages.count_followers
-
+    num_followers, num_following = g.user.count_followers()
     collection = {
         '@context': 'https://www.w3.org/ns/activitystreams',
         'id': request.base_url,
         'type': 'Collection',
         'summary': f"{id}'s {collection}",
-        'totalItems': count,
+        'totalItems': num_followers if collection == 'followers' else num_following,
         'first': page,
     }
     logger.info(f'Returning {json_dumps(collection, indent=2)}')

models.py

@@ -495,6 +495,21 @@ class User(StringIdModel, metaclass=ProtocolUserMeta):
   <img src="{img}" class="profile">
   {self.name()}</a>"""
 
+    # TODO: cache?
+    def count_followers(self):
+        """Counts this user's followers and followings.
+
+        Returns:
+          (int, int) tuple: (number of followers, number following)
+        """
+        num_followers = Follower.query(Follower.to == self.key,
+                                       Follower.status == 'active')\
+                                .count()
+        num_following = Follower.query(Follower.from_ == self.key,
+                                       Follower.status == 'active')\
+                                .count()
+        return num_followers, num_following
+
 
 class Object(StringIdModel):
     """An activity or other object, eg actor.

pages.py

@@ -121,7 +121,7 @@ def profile(protocol, id):
     load_user(protocol, id)
     query = Object.query(Object.users == g.user.key)
     objects, before, after = fetch_objects(query, by=Object.updated)
-    num_followers, num_following = count_followers()
+    num_followers, num_following = g.user.count_followers()
     return render_template('profile.html', **TEMPLATE_VARS, **locals())
 
 
@@ -161,7 +161,7 @@ def followers_or_following(protocol, id, collection):
     load_user(protocol, id)
 
     followers, before, after = Follower.fetch_page(collection)
-    num_followers, num_following = count_followers()
+    num_followers, num_following = g.user.count_followers()
     return render_template(
         f'{collection}.html',
         address=request.args.get('address'),
@@ -172,22 +172,6 @@ def followers_or_following(protocol, id, collection):
     )
 
 
-# TODO: cache?
-def count_followers():
-    start = time.time()
-    num_followers = Follower.query(Follower.to == g.user.key,
-                                   Follower.status == 'active')\
-                            .count()
-    end = time.time()
-    logger.info(f"Loading {g.user.key.id()}'s followers took {end - start}s")
-
-    num_following = Follower.query(Follower.from_ == g.user.key,
-                                   Follower.status == 'active')\
-                            .count()
-
-    return num_followers, num_following
-
-
 @app.get(f'/<any({",".join(PROTOCOLS)}):protocol>/<id>/feed')
 @canonicalize_request_domain(common.PROTOCOL_DOMAINS, common.PRIMARY_DOMAIN)
 def feed(protocol, id):

tests/test_models.py

@@ -226,6 +226,17 @@ class UserTest(TestCase):
         user.copies = [Target(uri='other:foo', protocol='other')]
         self.assertEqual('other:foo', user.get_copy(OtherFake))
 
+    def test_count_followers(self):
+        self.assertEqual((0, 0), g.user.count_followers())
+
+        Follower(from_=g.user.key, to=Fake(id='a').key).put()
+        Follower(from_=g.user.key, to=Fake(id='b').key).put()
+        Follower(from_=Fake(id='c').key, to=g.user.key).put()
+
+        del g.user
+        user = Web.get_by_id('y.z')
+        self.assertEqual((1, 2), user.count_followers())
+
 
 class ObjectTest(TestCase):
     def setUp(self):