Make it async

Author: fspv

generate.py

@@ -1,11 +1,13 @@
 #!/usr/bin/env python3
+
 import argparse
+import asyncio
 import json
 import logging
 import os
 import time
 from functools import lru_cache
-from typing import Dict, Iterator, List, Tuple
+from typing import Any, Coroutine, Dict, Iterator, List, Tuple
 
 import diskcache
 # https://github.com/kerrickstaley/genanki
@@ -25,6 +27,8 @@
 CACHE_DIR = "cache"
 ALLOWED_EXTENSIONS = {".py", ".go"}
 
+leetcode_api_access_lock = asyncio.Lock()
+
 
 logging.getLogger().setLevel(logging.INFO)
 
@@ -66,12 +70,10 @@ def __init__(self) -> None:
             os.mkdir(CACHE_DIR)
         self._cache = diskcache.Cache(CACHE_DIR)
 
-    def _get_problem_data(self, problem_slug: str) -> Dict[str, str]:
+    async def _get_problem_data(self, problem_slug: str) -> Dict[str, str]:
         if problem_slug in self._cache:
             return self._cache[problem_slug]
 
-        time.sleep(2)  # Leetcode has a rate limiter
-
         api_instance = self._api_instance
 
         graphql_request = leetcode.GraphqlQuery(
@@ -134,35 +136,40 @@ def _get_problem_data(self, problem_slug: str) -> Dict[str, str]:
             variables=leetcode.GraphqlQueryVariables(title_slug=problem_slug),
             operation_name="getQuestionDetail",
         )
-        data = api_instance.graphql_post(body=graphql_request).data.question
+
+        # Critical section. Don't allow more than one parallel request to
+        # the Leetcode API
+        async with leetcode_api_access_lock:
+            time.sleep(2)  # Leetcode has a rate limiter
+            data = api_instance.graphql_post(body=graphql_request).data.question
 
         # Save data in the cache
         self._cache[problem_slug] = data
 
         return data
 
-    def _get_description(self, problem_slug: str) -> str:
-        data = self._get_problem_data(problem_slug)
+    async def _get_description(self, problem_slug: str) -> str:
+        data = await self._get_problem_data(problem_slug)
         return data.content or "No content"
 
-    def _stats(self, problem_slug: str) -> Dict[str, str]:
-        data = self._get_problem_data(problem_slug)
+    async def _stats(self, problem_slug: str) -> Dict[str, str]:
+        data = await self._get_problem_data(problem_slug)
         return json.loads(data.stats)
 
-    def submissions_total(self, problem_slug: str) -> int:
-        return self._stats(problem_slug)["totalSubmissionRaw"]
+    async def submissions_total(self, problem_slug: str) -> int:
+        return (await self._stats(problem_slug))["totalSubmissionRaw"]
 
-    def submissions_accepted(self, problem_slug: str) -> int:
-        return self._stats(problem_slug)["totalAcceptedRaw"]
+    async def submissions_accepted(self, problem_slug: str) -> int:
+        return (await self._stats(problem_slug))["totalAcceptedRaw"]
 
-    def description(self, problem_slug: str) -> str:
-        return self._get_description(problem_slug)
+    async def description(self, problem_slug: str) -> str:
+        return await self._get_description(problem_slug)
 
-    def solution(self, problem_slug: str) -> str:
+    async def solution(self, problem_slug: str) -> str:
         return ""
 
-    def difficulty(self, problem_slug: str) -> str:
-        data = self._get_problem_data(problem_slug)
+    async def difficulty(self, problem_slug: str) -> str:
+        data = await self._get_problem_data(problem_slug)
         diff = data.difficulty
 
         if diff == "Easy":
@@ -174,34 +181,34 @@ def difficulty(self, problem_slug: str) -> str:
         else:
             raise ValueError(f"Incorrect difficulty: {diff}")
 
-    def paid(self, problem_slug: str) -> str:
-        data = self._get_problem_data(problem_slug)
+    async def paid(self, problem_slug: str) -> str:
+        data = await self._get_problem_data(problem_slug)
         return data.is_paid_only
 
-    def problem_id(self, problem_slug: str) -> str:
-        data = self._get_problem_data(problem_slug)
+    async def problem_id(self, problem_slug: str) -> str:
+        data = await self._get_problem_data(problem_slug)
         return data.question_frontend_id
 
-    def likes(self, problem_slug: str) -> int:
-        data = self._get_problem_data(problem_slug)
+    async def likes(self, problem_slug: str) -> int:
+        data = await self._get_problem_data(problem_slug)
         likes = data.likes
 
         if not isinstance(likes, int):
             raise ValueError(f"Likes should be int: {likes}")
 
         return likes
 
-    def dislikes(self, problem_slug: str) -> int:
-        data = self._get_problem_data(problem_slug)
+    async def dislikes(self, problem_slug: str) -> int:
+        data = await self._get_problem_data(problem_slug)
         dislikes = data.dislikes
 
         if not isinstance(dislikes, int):
             raise ValueError(f"Dislikes should be int: {dislikes}")
 
         return dislikes
 
-    def tags(self, problem_slug: str) -> List[str]:
-        data = self._get_problem_data(problem_slug)
+    async def tags(self, problem_slug: str) -> List[str]:
+        data = await self._get_problem_data(problem_slug)
         return list(map(lambda x: x.slug, data.topic_tags))
 
 
@@ -237,7 +244,7 @@ def get_leetcode_task_handles() -> Iterator[Tuple[str, str, str]]:
             yield (topic, stat.question__title, stat.question__title_slug)
 
 
-def generate_anki_note(
+async def generate_anki_note(
     leetcode_data: LeetcodeData,
     leetcode_model: genanki.Model,
     leetcode_task_handle: str,
@@ -248,29 +255,29 @@ def generate_anki_note(
         model=leetcode_model,
         fields=[
             leetcode_task_handle,
-            str(leetcode_data.problem_id(leetcode_task_handle)),
+            str(await leetcode_data.problem_id(leetcode_task_handle)),
             leetcode_task_title,
             topic,
-            leetcode_data.description(leetcode_task_handle),
-            leetcode_data.difficulty(leetcode_task_handle),
-            "yes" if leetcode_data.paid(leetcode_task_handle) else "no",
-            str(leetcode_data.likes(leetcode_task_handle)),
-            str(leetcode_data.dislikes(leetcode_task_handle)),
-            str(leetcode_data.submissions_total(leetcode_task_handle)),
-            str(leetcode_data.submissions_accepted(leetcode_task_handle)),
+            await leetcode_data.description(leetcode_task_handle),
+            await leetcode_data.difficulty(leetcode_task_handle),
+            "yes" if await leetcode_data.paid(leetcode_task_handle) else "no",
+            str(await leetcode_data.likes(leetcode_task_handle)),
+            str(await leetcode_data.dislikes(leetcode_task_handle)),
+            str(await leetcode_data.submissions_total(leetcode_task_handle)),
+            str(await leetcode_data.submissions_accepted(leetcode_task_handle)),
             str(
                 int(
-                    leetcode_data.submissions_accepted(leetcode_task_handle)
-                    / leetcode_data.submissions_total(leetcode_task_handle)
+                    await leetcode_data.submissions_accepted(leetcode_task_handle)
+                    / await leetcode_data.submissions_total(leetcode_task_handle)
                     * 100
                 )
             ),
         ],
-        tags=leetcode_data.tags(leetcode_task_handle),
+        tags=await leetcode_data.tags(leetcode_task_handle),
     )
 
 
-def generate(start: int, stop: int) -> None:
+async def generate(start: int, stop: int) -> None:
     leetcode_model = genanki.Model(
         LEETCODE_ANKI_MODEL_ID,
         "Leetcode model",
@@ -328,27 +335,35 @@ def generate(start: int, stop: int) -> None:
     )
     leetcode_deck = genanki.Deck(LEETCODE_ANKI_DECK_ID, "leetcode")
     leetcode_data = LeetcodeData()
-    for topic, leetcode_task_title, leetcode_task_handle in tqdm(
-        list(get_leetcode_task_handles())[start:stop]
-    ):
-        leetcode_note = generate_anki_note(
-            leetcode_data,
-            leetcode_model,
-            leetcode_task_handle,
-            leetcode_task_title,
-            topic,
+
+    note_generators: List[Coroutine[Any, Any, LeetcodeNote]] = []
+
+    for topic, leetcode_task_title, leetcode_task_handle in list(
+        get_leetcode_task_handles()
+    )[start:stop]:
+        note_generators.append(
+            generate_anki_note(
+                leetcode_data,
+                leetcode_model,
+                leetcode_task_handle,
+                leetcode_task_title,
+                topic,
+            )
         )
-        leetcode_deck.add_note(leetcode_note)
+
+    for leetcode_note in tqdm(note_generators):
+        leetcode_deck.add_note(await leetcode_note)
 
     genanki.Package(leetcode_deck).write_to_file(OUTPUT_FILE)
 
 
-def main() -> None:
+async def main() -> None:
     args = parse_args()
 
     start, stop = args.start, args.stop
-    generate(start, stop)
+    await generate(start, stop)
 
 
 if __name__ == "__main__":
-    main()
+    loop = asyncio.get_event_loop()
+    loop.run_until_complete(main())