fix: serialize None according to return type (#266)

Fixes #265

lang/python.go

@@ -117,12 +117,10 @@ func toPythonType(typeName string) string {
 }
 
 func (p python) generateNormalTestCode(q *leetcode.QuestionData) (string, error) {
-	const template = `if __name__ == "__main__":
-%s
-	print("\n%s", serialize(ans))
-`
-	code := ""
+	code := "if __name__ == \"__main__\":\n"
+
 	paramNames := make([]string, 0, len(q.MetaData.Params))
+	paramTypes := make([]string, 0, len(q.MetaData.Params))
 	for _, param := range q.MetaData.Params {
 		varName := param.Name
 		varType := toPythonType(param.Type)
@@ -133,13 +131,15 @@ func (p python) generateNormalTestCode(q *leetcode.QuestionData) (string, error)
 			varType,
 		)
 		paramNames = append(paramNames, param.Name)
+		paramTypes = append(paramTypes, varType)
 	}
 	if q.MetaData.Return != nil && q.MetaData.Return.Type != "void" {
 		code += fmt.Sprintf(
 			"\tans = Solution().%s(%s)\n",
 			q.MetaData.Name,
 			strings.Join(paramNames, ", "),
 		)
+		code += fmt.Sprintf("\tprint(\"\\n%s\", serialize(ans, \"%s\"))\n", testCaseOutputMark, q.MetaData.Return.Type)
 	} else {
 		code += fmt.Sprintf(
 			"\t%s(%s)\n",
@@ -149,16 +149,21 @@ func (p python) generateNormalTestCode(q *leetcode.QuestionData) (string, error)
 		if q.MetaData.Output != nil {
 			ansName := paramNames[q.MetaData.Output.ParamIndex]
 			code += fmt.Sprintf("\tans = %s\n", ansName)
+			code += fmt.Sprintf(
+				"\tprint(\"\\n%s\", serialize(ans, \"%s\"))\n",
+				testCaseOutputMark,
+				paramTypes[q.MetaData.Output.ParamIndex],
+			)
 		} else {
 			code += "\tans = None\n"
+			code += fmt.Sprintf("\tprint(\"\\n%s\", \"null\")\n", testCaseOutputMark)
 		}
 	}
 
-	testContent := fmt.Sprintf(template, code, testCaseOutputMark)
 	if q.MetaData.Manual {
-		testContent = fmt.Sprintf("# %s\n%s", manualWarning, testContent)
+		code = fmt.Sprintf("# %s\n%s", manualWarning, code)
 	}
-	return testContent, nil
+	return code, nil
 }
 
 func (p python) generateSystemDesignTestCode(q *leetcode.QuestionData) (string, error) {

testutils/python/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "leetgo-py"
-version = "0.2.3"
+version = "0.2.4"
 description = "Python test utils for leetgo"
 authors = ["j178 <[email protected]>"]
 license = "MIT"

testutils/python/src/leetgo_py/parse.py

@@ -14,8 +14,14 @@ def split_array(s: str) -> List[str]:
     return res
 
 
-def serialize(val: Any) -> str:
-    if isinstance(val, bool):
+def serialize(val: Any, ty: str = None) -> str:
+    if val is None:
+        if ty is None:
+            raise Exception("None value without type")
+        if ty == "ListNode" or ty == "TreeNode":
+            return "[]"
+        return "null"
+    elif isinstance(val, bool):
         return "true" if val else "false"
     elif isinstance(val, int):
         return str(val)