Change LeetCode get_submission method into a generator to avoid API rate limit problems

Author: NeverMendel

README.md

@@ -25,10 +25,11 @@ section [Docker Image here](#docker-image).
 
 To use `leetcode-export` you can either download it from pypi.org or you can clone this repository.
 
-## Download from pypi.org
+### Download from pypi.org
 
-Run `pip install leetcode-export` to download it and install all the needed dependencies. You might need to use `pip3`
-depending on the configuration of your system.
+Run `pip install leetcode-export` to install leetcode-export, you might have to use `pip3` of your system. To use the
+script run `leetcode-export`, optionally supply the script arguments, for more instructions read the
+section [script arguments here](#script-arguments).
 
 ### Clone the repository
 
@@ -44,14 +45,14 @@ Install all the needed dependencies:
 pip install -r requirements.txt
 ```
 
-Now you can either install it or just execute it:
+You can either install leetcode-export in your system or just execute it:
 
-- To install leetcode-export in your system:
+- To install it run:
     ```bash
     python setup.py install
     ```
 
-- To run the project without installing it:
+- To execute the script without installing it:
     ```bash
     python -m leetcode_export --folder submissions
     ```
@@ -70,6 +71,25 @@ Download all your LeetCode submission in the current folder:
 docker run -it -v $(pwd):/usr/app/out --rm nevermendel/leetcode-export
 ```
 
+## Login
+
+To download your submissions you need to log in your LeetCode account by providing the cookies. To log in using cookies,
+you need to get them from a session where you are already logged in.
+
+**Steps required**:
+
+- Login in your LeetCode account in your browser
+- Open the browser's Dev Tool
+- Click on the Network tab
+- Copy the cookie header that can be found under Request Headers in any leetcode.com request.
+
+You can insert the cookie string that you have just copied in the interactive menu (recommended) or you can pass it as a
+program argument when lunching the script, like in the following example:
+
+```bash
+python leetcode-export --cookies {COOKIES}
+```
+
 ## Script arguments
 
 The script accepts the following arguments:
@@ -93,24 +113,6 @@ optional arguments:
                         Submission filename format
 ```
 
-## Login
-
-To download your submissions you need to log in your LeetCode account by providing the cookies. To log in using cookies,
-you need to get them from a session where you are already logged in.
-
-**Steps required**: 
-- Login in your LeetCode account in your browser
-- Open the browser's Dev Tool
-- Click on the Network tab
-- Copy the cookie header that can be found under Request Headers in any leetcode.com request.
-
-You can insert the cookie string that you have just copied in the interactive menu (recommended) or you can pass it as a
-program argument when lunching the script, like in the following example:
-
-```bash
-python leetcode-export --cookies {COOKIES}
-```
-
 Using the interactive menu is preferred as it will avoid storing your cookies in the command history.
 
 ## Filename template arguments
@@ -166,7 +168,6 @@ date_formatted: str
 extension: str
 ```
 
-
 Default submission filename
 template: `${date_formatted} - ${status_display} - runtime ${runtime} - memory ${memory}.${extension}`
 

leetcode_export/__main__.py

@@ -3,10 +3,9 @@
 import logging
 import os
 from string import Template
-from typing import List, Dict
+from typing import Set
 
 from leetcode_export.leetcode import LeetCode
-from leetcode_export.leetcode_rest import Submission
 
 PROBLEM_CONTENT_TEMPLATE = Template('''${question_id} - ${title}
 ${difficulty} - https://leetcode.com/problems/${title_slug}/
@@ -70,31 +69,30 @@ def main():
         os.mkdir(args.folder)
     os.chdir(args.folder)
 
-    submissions: Dict[str, List[Submission]] = leetcode.get_submissions()
-
-    for slug in submissions:
-        logging.info(f"Processing {slug}")
-        if not os.path.exists(slug):
-            os.mkdir(slug)
+    written_problems: Set[str] = set()
+
+    for submission in leetcode.get_submissions():
+        if not os.path.exists(submission.title_slug):
+            os.mkdir(submission.title_slug)
+        os.chdir(submission.title_slug)
+
+        if submission.title_slug not in written_problems:
+            problem = leetcode.get_problem(submission.title_slug)
+            info_filename = problem_template.substitute(**problem.__dict__)
+            if not os.path.exists(info_filename):
+                info_file = open(info_filename, 'w+')
+                info_file.write(PROBLEM_CONTENT_TEMPLATE.substitute(**problem.__dict__))
+                info_file.close()
+            written_problems.add(submission.title_slug)
+
+        sub_filename = submission_template.substitute(**submission.__dict__)
+        if not os.path.exists(sub_filename):
+            logging.info(f"Writing {submission.title_slug}/{sub_filename}")
+            sub_file = open(sub_filename, 'w+')
+            sub_file.write(submission.code)
+            sub_file.close()
         else:
-            logging.info(f"Folder {slug} already exists")
-        os.chdir(slug)
-        problem_description = leetcode.get_problem(slug)
-        info_filename = problem_template.substitute(**problem_description.__dict__)
-        if not os.path.exists(info_filename):
-            info_file = open(info_filename, 'w+')
-            info_file.write(PROBLEM_CONTENT_TEMPLATE.substitute(**problem_description.__dict__))
-            info_file.close()
-
-        for sub in submissions[slug]:
-            sub_filename = submission_template.substitute(**sub.__dict__)
-            if not os.path.exists(sub_filename):
-                logging.info(f"Writing {slug}/{sub_filename}")
-                sub_file = open(sub_filename, 'w+')
-                sub_file.write(sub.code)
-                sub_file.close()
-            else:
-                logging.info(f"{slug}/{sub_filename} already exists, skipping it")
+            logging.info(f"{submission.title_slug}/{sub_filename} already exists, skipping it")
 
         os.chdir("..")
 

leetcode_export/leetcode.py

@@ -1,7 +1,7 @@
 import datetime
 import logging
 from time import sleep
-from typing import Dict, List
+from typing import Dict, Iterator
 
 import requests
 
@@ -85,15 +85,15 @@ def get_problem(self, slug: str) -> Problem:
             problem_dict = dict_camelcase_to_snakecase(response.json()['data']['question'])
             return Problem.from_dict(problem_dict)
 
-    def get_submissions(self) -> Dict[str, List[Submission]]:
+    def get_submissions(self) -> Iterator[Submission]:
         '''
-        Get list of submission for logged user
-        :return: Dict[str, List[Submission]], dictionary with slug as key and submission for given problem as value
+        Get submissions for logged user
+        :return: Iterator[Submission], LeetCode submission
         '''
         if not self.is_user_logged():
             logging.warning("Trying to get user submissions while user is not logged in")
-            return {}
-        dictionary: Dict[str, List[Submission]] = {}
+            return None
+
         current = 0
         response_json: Dict = {'has_next': True}
         while 'detail' not in response_json and 'has_next' in response_json and response_json['has_next']:
@@ -112,12 +112,9 @@ def get_submissions(self) -> Dict[str, List[Submission]]:
                         if type(submission_dict[key]) == str and key != 'url' and key != 'code':
                             submission_dict[key] = remove_special_characters(submission_dict[key])
                     submission = Submission.from_dict(submission_dict)
-                    if submission.title_slug not in dictionary:
-                        dictionary[submission.title_slug] = []
-                    dictionary[submission.title_slug].append(submission)
+                    yield submission
 
             current += 20
             sleep(1)
         if 'detail' in response_json:
             logging.warning(response_json['detail'])
-        return dictionary

setup.py

@@ -5,7 +5,7 @@
 
 setup(
     name='leetcode-export',
-    version='1.1',
+    version='1.1.1',
     url='https://github.com/NeverMendel/leetcode-export',
     license='MIT',
     author='Davide Cazzin',