Fix smart word-wrap for "sbase print FILE.py"

Author: mdmintz

seleniumbase/console_scripts/sb_print.py

@@ -198,11 +198,15 @@ def main():
                         new_sb_lines.append(line1)
                         new_sb_lines.append(line2)
                         continue
-                    if line.count("(") == 1 and line.count(")") == 1:
+                    if (
+                        line.count("(") >= 1
+                        and line.count("(") == line.count(")")
+                    ):
                         whitespace = line_length2 - len(line.lstrip())
                         new_ws = line[0:whitespace] + "    "
-                        line1 = line.split("(")[0] + "("
-                        line2 = new_ws + line.split("(")[1]
+                        first_paren = line.find("(")
+                        line1 = line[:first_paren + 1]
+                        line2 = new_ws + line[first_paren + 1:]
                         if not ("):") in line2:
                             new_sb_lines.append(line1)
                             if get_width(line2) + w > console_width:
@@ -302,9 +306,7 @@ def main():
                                     continue
                                 elif line2.count(", ") == 1:
                                     line2a = line2.split(", ")[0] + ","
-                                    line2b = new_ws + (
-                                        line2.split(", ")[1]
-                                    )
+                                    line2b = new_ws + line2.split(", ")[1]
                                     new_sb_lines.append(line2a)
                                     new_sb_lines.append(line2b)
                                     continue

seleniumbase/translate/translator.py

@@ -620,11 +620,15 @@ def main():
                         new_sb_lines.append(line1)
                         new_sb_lines.append(line2)
                         continue
-                    if line.count("(") == 1 and line.count(")") == 1:
+                    if (
+                        line.count("(") >= 1
+                        and line.count("(") == line.count(")")
+                    ):
                         whitespace = line_length2 - len(line.lstrip())
                         new_ws = line[0:whitespace] + "    "
-                        line1 = line.split("(")[0] + "("
-                        line2 = new_ws + line.split("(")[1]
+                        first_paren = line.find("(")
+                        line1 = line[:first_paren + 1]
+                        line2 = new_ws + line[first_paren + 1:]
                         if not ("):") in line2:
                             new_sb_lines.append(line1)
                             if get_width(line2) + w > console_width:
@@ -722,9 +726,7 @@ def main():
                                     continue
                                 elif line2.count(", ") == 1:
                                     line2a = line2.split(", ")[0] + ","
-                                    line2b = new_ws + (
-                                        line2.split(", ")[1]
-                                    )
+                                    line2b = new_ws + line2.split(", ")[1]
                                     new_sb_lines.append(line2a)
                                     new_sb_lines.append(line2b)
                                     continue