feat: update main script

Author: yanglbme

solution/main.py

@@ -4,6 +4,33 @@
 
 from spider import Spider
 
+# load templates
+with open('./readme_template.md', 'r', encoding='utf-8') as f:
+    readme_cn = f.read()
+with open('./readme_template_en.md', 'r', encoding='utf-8') as f:
+    readme_en = f.read()
+
+with open('./problem_readme_template.md', 'r', encoding='utf-8') as f:
+    problem_readme_cn = f.read()
+with open('./problem_readme_template_en.md', 'r', encoding='utf-8') as f:
+    problem_readme_en = f.read()
+with open('./sql_problem_readme_template.md', 'r', encoding='utf-8') as f:
+    sql_readme_cn = f.read()
+with open('./sql_problem_readme_template_en.md', 'r', encoding='utf-8') as f:
+    sql_readme_en = f.read()
+with open('./bash_problem_readme_template.md', 'r', encoding='utf-8') as f:
+    bash_readme_cn = f.read()
+with open('./bash_problem_readme_template_en.md', 'r', encoding='utf-8') as f:
+    bash_readme_en = f.read()
+
+
+def select_templates(category):
+    if category == 'Shell':
+        return [bash_readme_cn, bash_readme_en]
+    if category == 'Database':
+        return [sql_readme_cn, sql_readme_en]
+    return [problem_readme_cn, problem_readme_en]
+
 
 def generate_readme(result):
     md_table_cn = [item['md_table_row_cn'] for item in result]
@@ -18,11 +45,6 @@ def generate_readme(result):
         )
     table_cn += ''.join(items)
 
-    with open('./readme_template.md', 'r', encoding='utf-8') as f:
-        readme_cn = f.read()
-    with open('./README.md', 'w', encoding='utf-8') as f:
-        f.write(readme_cn.format(table_cn))
-
     # generate README_EN.md
     items = []
     table_en = '\n|  #  |  Solution  |  Tags  |  Difficulty  |  Remark |\n| --- | --- | --- | --- | --- |'
@@ -32,26 +54,13 @@ def generate_readme(result):
         )
     table_en += ''.join(items)
 
-    with open('./readme_template_en.md', 'r', encoding='utf-8') as f:
-        readme_en = f.read()
+    with open('./README.md', 'w', encoding='utf-8') as f:
+        f.write(readme_cn.format(table_cn))
     with open('./README_EN.md', 'w', encoding='utf-8') as f:
         f.write(readme_en.format(table_en))
 
 
 def generate_question_readme(result):
-    with open('./problem_readme_template.md', 'r', encoding='utf-8') as f:
-        readme_cn = f.read()
-    with open('./problem_readme_template_en.md', 'r', encoding='utf-8') as f:
-        readme_en = f.read()
-    with open('./sql_problem_readme_template.md', 'r', encoding='utf-8') as f:
-        sql_readme_cn = f.read()
-    with open('./sql_problem_readme_template_en.md', 'r', encoding='utf-8') as f:
-        sql_readme_en = f.read()
-    with open('./bash_problem_readme_template.md', 'r', encoding='utf-8') as f:
-        bash_readme_cn = f.read()
-    with open('./bash_problem_readme_template_en.md', 'r', encoding='utf-8') as f:
-        bash_readme_en = f.read()
-
     for item in result:
         if not item['content_cn'] and not item['content_en']:
             continue
@@ -62,17 +71,10 @@ def generate_question_readme(result):
         if os.path.isdir(path):
             continue
         os.makedirs(path)
+
         # choose the readme template
         category = item['category']
-        if category == 'Shell':
-            readme_template_cn = bash_readme_cn
-            readme_template_en = bash_readme_en
-        elif category == 'Database':
-            readme_template_cn = sql_readme_cn
-            readme_template_en = sql_readme_en
-        else:
-            readme_template_cn = readme_cn
-            readme_template_en = readme_en
+        readme_template_cn, readme_template_en = select_templates(category)
 
         # generate lc-cn problem readme
         with open(f'{path}/README.md', 'w', encoding='utf-8') as f1:
@@ -101,8 +103,7 @@ def generate_question_readme(result):
 
 def generate_summary(result):
     """generate summary files"""
-    summary_cn = ''
-    summary_en = ''
+    summary_cn = summary_en = ''
     m = {int(item['frontend_question_id']): item for item in result}
     for file in os.listdir('./'):
         if os.path.isdir("./" + file) and file != '__pycache__':
@@ -135,10 +136,15 @@ def save(result):
 
 
 if __name__ == '__main__':
-    cn = ''
-    en = ''
-    spider = Spider(cn, en)
+    cookie_cn = ''
+    cookie_en = ''
+    spider = Spider(cookie_cn, cookie_en)
     res = spider.run()
+
+    # with open('./result.json', 'r', encoding='utf-8') as f:
+    #     res = f.read()
+    #     res = json.loads(res)
+
     generate_readme(res)
     generate_question_readme(res)
     generate_summary(res)

solution/spider.py

@@ -19,6 +19,71 @@
 cn_graph_url = 'https://leetcode-cn.com/graphql'
 difficulty = dict(Easy='简单', Medium='中等', Hard='困难')
 
+"""
+{
+    "sub_folder":"2200-2299",
+    "question_id":"2385",
+    "frontend_question_id":"2237",
+    "paid_only":true,
+    "paid_only_cn":true,
+    "category":"Algorithms",
+    "url_cn":"https://leetcode-cn.com/problems/count-positions-on-street-with-required-brightness",
+    "url_en":"https://leetcode.com/problems/count-positions-on-street-with-required-brightness",
+    "relative_path_cn":"/solution/2200-2299/2237.Count%20Positions%20on%20Street%20With%20Required%20Brightness/README.md",
+    "relative_path_en":"/solution/2200-2299/2237.Count%20Positions%20on%20Street%20With%20Required%20Brightness/README_EN.md",
+    "title_cn":"Count Positions on Street With Required Brightness",
+    "title_en":"Count Positions on Street With Required Brightness",
+    "question_title_slug":"count-positions-on-street-with-required-brightness",
+    "content_en":"Description",
+    "content_cn":"问题描述[此处省略]",
+    "tags_en":[
+
+    ],
+    "tags_cn":[
+
+    ],
+    "difficulty_en":"Medium",
+    "difficulty_cn":"中等",
+    "code_snippets":[
+        {
+            "lang":"C++",
+            "langSlug":"cpp",
+            "code":"class Solution {
+public:
+    int meetRequirement(int n, vector<vector<int>>& lights, vector<int>& requirement) {
+        
+    }
+};",
+            "__typename":"CodeSnippetNode"
+        },
+        {
+            "lang":"Java",
+            "langSlug":"java",
+            "code":"class Solution {
+    public int meetRequirement(int n, int[][] lights, int[] requirement) {
+
+    }
+}",
+            "__typename":"CodeSnippetNode"
+        }
+    ],
+    "md_table_row_cn":[
+        "[2237](https://leetcode-cn.com/problems/count-positions-on-street-with-required-brightness)",
+        "[Count Positions on Street With Required Brightness](/solution/2200-2299/2237.Count%20Positions%20on%20Street%20With%20Required%20Brightness/README.md)",
+        "",
+        "中等",
+        "🔒"
+    ],
+    "md_table_row_en":[
+        "[2237](https://leetcode.com/problems/count-positions-on-street-with-required-brightness)",
+        "[Count Positions on Street With Required Brightness](/solution/2200-2299/2237.Count%20Positions%20on%20Street%20With%20Required%20Brightness/README_EN.md)",
+        "",
+        "Medium",
+        "🔒"
+    ]
+}
+"""
+
 
 class Spider:
     def __init__(self, cookie_cn: str, cookie_en: str):