Fix bugs in crawler and Python codegen

- Crawler also needs to catch "NoSuchElementException" when checking
  whether the problem page contains an interactive IDE.
- Python codegen needs to import typing and define TreeNode before the
  Solution class.

Author: huzecong

lchelper/codegen/python.py

@@ -27,6 +27,14 @@ def line_comment_symbol(self) -> str:
     @property
     def template_code(self) -> str:
         return r"""
+from typing import *
+
+class TreeNode:
+     def __init__(self, x):
+         self.val = x
+         self.left = None
+         self.right = None
+
 # BEGIN SUBMIT
 
 # BEGIN USER TEMPLATE
@@ -43,14 +51,6 @@ def template_code(self) -> str:
 
 # END STATEMENT
 
-from typing import List, Optional
-
-class TreeNode:
-     def __init__(self, x):
-         self.val = x
-         self.left = None
-         self.right = None
-
 
 def _construct_tree(parent: List[Optional[int]], idx: int = 0) -> Optional[TreeNode]:
     if idx >= len(parent) or parent[idx] is None:
@@ -79,11 +79,11 @@ def test(msg: str, a, b):
         "string": "str",
     }
 
-    def _convert_cpp_type_to_py(self, type_name: str) -> str:
+    def _convert_cpp_type(self, type_name: str) -> str:
         type_name = type_name.strip().rstrip("*&")
         if type_name.startswith("vector<") and type_name.endswith(">"):
             inner_type_name = type_name[len("vector<"):-len(">")]
-            return f"List[{self._convert_cpp_type_to_py(inner_type_name)}]"
+            return f"List[{self._convert_cpp_type(inner_type_name)}]"
         return self.TYPE_MAP.get(type_name, type_name)
 
     def generate_solution_code(self, signature: Signature) -> Code:
@@ -95,15 +95,15 @@ def generate_solution_code(self, signature: Signature) -> Code:
             functions = [signature.function]
         fn_codes = []
         for func_sig in functions:
-            args = ", ".join(f"{arg_name}: {self._convert_cpp_type_to_py(arg_type)}"
+            args = ", ".join(f"{arg_name}: {self._convert_cpp_type(arg_type)}"
                              for arg_type, arg_name in func_sig.arguments)
             if func_sig.name == class_name:
                 fn_code = [
                     f"    def __init__(self, {args}):",
                     f"        pass"]
             else:
                 fn_code = [
-                    f"    def {func_sig.name}(self, {args}) -> {self._convert_cpp_type_to_py(func_sig.return_type)}:",
+                    f"    def {func_sig.name}(self, {args}) -> {self._convert_cpp_type(func_sig.return_type)}:",
                     f"        pass"]
             fn_codes.append(fn_code)
         code = [f"class {class_name}:"] + self.list_join(fn_codes, [""])
@@ -128,7 +128,7 @@ def to_tree(parent: List[Optional[int]]) -> str:
             return f"_construct_tree([{', '.join('None' if x is None else str(x) for x in parent)}])"
 
         def to_val(val: Any, type_name: str) -> str:
-            if type_name == "TreeNode":
+            if self._convert_cpp_type(type_name) == "TreeNode":
                 return to_tree(val)
             return to_str(val)
 

lchelper/crawler.py

@@ -3,7 +3,7 @@
 from typing import List
 
 from selenium import webdriver
-from selenium.common.exceptions import TimeoutException
+from selenium.common.exceptions import NoSuchElementException, TimeoutException
 from selenium.webdriver.common.by import By
 from selenium.webdriver.support import expected_conditions as Expected
 from selenium.webdriver.support.wait import WebDriverWait
@@ -157,7 +157,7 @@ def get_problems(contest_url: str, site: str, cookie_path: str) -> List[Problem]
             statement_css_selector = "div.question-content"
             code_css_selector = "pre.CodeMirror-line"
             statement = browser.find_element_by_css_selector(statement_css_selector).text
-        except TimeoutException:
+        except (TimeoutException, NoSuchElementException):
             # Page after contest; statement and editor in vertically split panes.
             statement_css_selector = "div[data-key='description-content'] div.content__1Y2H"
             code_css_selector = "div.monaco-scrollable-element div.view-line"