Improve and refactor "sbase print FILE"

Author: mdmintz

seleniumbase/console_scripts/rich_helper.py

@@ -0,0 +1,46 @@
+from rich.console import Console
+from rich.markdown import Markdown
+from rich.syntax import Syntax
+
+
+def process_syntax(code, lang, theme, line_numbers, code_width, word_wrap):
+    syntax = Syntax(
+        code,
+        lang,
+        theme=theme,
+        line_numbers=line_numbers,
+        code_width=code_width,
+        word_wrap=word_wrap,
+    )
+    return syntax
+
+
+def display_markdown(code):
+    try:
+        markdown = Markdown(code)
+        console = Console()
+        console.print(markdown)  # noqa
+        return True  # Success
+    except Exception:
+        return False  # Failure
+
+
+def display_code(code):
+    try:
+        console = Console()
+        console.print(code)  # noqa
+        return True  # Success
+    except Exception:
+        return False  # Failure
+
+
+def fix_emoji_spacing(code):
+    try:
+        # Fix the display width of certain emojis that take up two spaces
+        double_width_emojis = ["🗺️", "🖼️", "🗄️", "⏺️", "♻️", "🗂️", "🖥️"]
+        for emoji in double_width_emojis:
+            if emoji in code:
+                code = code.replace(emoji, emoji + " ")
+    except Exception:
+        pass
+    return code

seleniumbase/console_scripts/run.py

@@ -760,14 +760,6 @@ def main():
             show_convert_usage()
     elif command == "print":
         if len(command_args) >= 1:
-            if sys.version_info[0] == 2:
-                c5 = colorama.Fore.RED + colorama.Back.LIGHTYELLOW_EX
-                cr = colorama.Style.RESET_ALL
-                msg = '"sbase print" does NOT support Python 2! '
-                msg += 'Try using the Unix "cat" command instead!'
-                message = "\n" + c5 + msg + cr + "\n"
-                print("")
-                raise Exception(message)
             from seleniumbase.console_scripts import sb_print
 
             sb_print.main()

seleniumbase/console_scripts/sb_print.py

@@ -114,15 +114,19 @@ def main():
     if help_me:
         invalid_run_command(invalid_cmd)
 
-    with open(file_to_print, "r", encoding="utf-8") as f:
-        all_code = f.read()
+    all_code = None
+    if sys.version_info[0] >= 3:
+        with open(file_to_print, "r", encoding="utf-8") as f:
+            all_code = f.read()
+    else:
+        with open(file_to_print, "r") as f:
+            all_code = f.read()
     all_code = all_code.replace("\t", "    ")
     code_lines = all_code.split("\n")
 
     console_width = None  # width of console output when running script
     used_width = None  # code_width and few spaces on right for padding
-    magic_console = None
-    magic_syntax = None
+    magic_syntax = None  # the syntax generated by rich.syntax.Syntax()
     try:
         console_width = os.popen("stty size", "r").read().split()[1]
         if console_width:
@@ -135,8 +139,7 @@ def main():
         use_rich = True
 
     if use_rich:
-        from rich.console import Console
-        from rich.syntax import Syntax
+        from seleniumbase.console_scripts import rich_helper
 
         the_code = "\n".join(code_lines)
         code_width = 1
@@ -566,22 +569,16 @@ def main():
         if console_width and (code_width + extra_r_spaces < console_width):
             used_width = code_width + extra_r_spaces
 
-        try:
-            if "🗺️" in the_code:
-                # Fix width of an emoji
-                the_code = the_code.replace("🗺️", "🗺️ ")
-        except Exception:
-            pass
+        the_code = rich_helper.fix_emoji_spacing(the_code)
 
-        magic_syntax = Syntax(
+        magic_syntax = rich_helper.process_syntax(
             the_code,
             code_lang,
             theme="monokai",
             line_numbers=line_numbers,
             code_width=used_width,
             word_wrap=word_wrap,
         )
-        magic_console = Console()
     # ----------------------------------------
     dash_length = 62  # May change
     if used_width and used_width + w < console_width:
@@ -591,22 +588,35 @@ def main():
     dashes = "-" * dash_length
     print(dashes)
     print_success = False
-    if use_rich and code_lang == "markdown":
-        try:
-            from rich.markdown import Markdown
-
-            markdown = Markdown(all_code)
-            markdown_console = Console()
-            markdown_console.print(markdown)  # noqa
-            print_success = True
-        except Exception:
-            pass
+    if use_rich and code_lang == "markdown" and not line_numbers:
+        all_code = rich_helper.fix_emoji_spacing(all_code)
+        all_code = all_code.replace("<br />", "\n")
+        all_code = all_code.replace(
+            '<p align="center">', '\n')
+        all_code = all_code.replace(
+            '<p align="left">', '\n')
+        all_code = all_code.replace(
+            '<p align="right">', '\n')
+        all_code = all_code.replace("<p>", "\n")
+        all_code = all_code.replace("</p>", "\n")
+        if "<b>*" not in all_code and "*<b>" not in all_code:
+            if "</b>*" not in all_code and "*</b>" not in all_code:
+                all_code = all_code.replace("<b>", "**")
+                all_code = all_code.replace("</b>", "**")
+        if "<code>`" not in all_code and "`<code>" not in all_code:
+            if "</code>`" not in all_code and "`</code>" not in all_code:
+                all_code = all_code.replace("<code>", "``")
+                all_code = all_code.replace("</code>", "``")
+        # Display ALL <h> tags as an <h1> because the font size is fixed
+        all_code = all_code.replace("\n<h1>", "\n# ").replace("</h1>", "")
+        all_code = all_code.replace("\n<h2>", "\n# ").replace("</h2>", "")
+        all_code = all_code.replace("\n<h3>", "\n# ").replace("</h3>", "")
+        all_code = all_code.replace("\n<h4>", "\n# ").replace("</h4>", "")
+        print_success = rich_helper.display_markdown(all_code)
+        if all_code.endswith("\n"):
+            print()  # Because "rich" skips the last line if new-line
     elif use_rich and magic_syntax:
-        try:
-            magic_console.print(magic_syntax)  # noqa
-            print_success = True
-        except Exception:
-            pass
+        print_success = rich_helper.display_code(magic_syntax)
     if not use_rich or not magic_syntax or not print_success:
         for line in code_lines:
             print(line)