from pygments import highlight
from pygments.style import Style
from pygments import token
from pygments.lexers import PythonLexer
from pygments.formatters import HtmlFormatter
from bs4 import BeautifulSoup


class CustomStyle(Style):
    default_style = ""
    styles = {
        token.Text: "",
        token.Keyword: "bold #ff9800",
        token.Name: "",
        token.Operator: "",

        token.Name.Class: "#ffeb3b",
        token.Name.Function: "#be9ddb",
        token.Name.Builtin: "#ab47bc",
        token.Name.Attribute: "#4db6ac",
        token.Name.Constant: "italic #ef5350",
        token.Name.Decorator: "#bcaaa4",
        token.Name.Function.Magic: "#e040fb",
        token.Name.Label: "#ff4081",

        token.String: "#8bc34a",
        token.Number: "#42a5f5",
        token.Punctuation: "#ff9800",
        token.Comment: "italic #b0bec5",
        token.Generic: "",
        token.Escape: "#ff5722",
        token.String.Escape: "#ff5722",
        token.String.Regex: "#ffffff",
        token.Comment.Preproc: "#536dfe",
        token.Comment.Hashbang: "#536dfe",
        token.Error: "#f44336",
        token.Other: "",
        token.Whitespace: "",
    }


code = """
___all___ = ["hello_world"]

def interesting_function(func: Callable) -> Callable:
    \"\"\"Print the function's result when called.
    
    Lorem ipsum dolor sit amet, consectetur adipiscing elit. Nullam nec nulla
    nec nulla laoreet fermentum. Integer id lacus ut odio interdum ultricies.
    \"\"\"
    def wrapper(*args, **kwargs):
        result = func(*args, **kwargs)          # avoid double evaluation
        print(result)
        return result
    return wrapper

@interesting_function
def hello_world():
    return "Hello, world!\\a"                   # a nice bell

if __name__ == "__main__":
    hello_world()
"""

lexer = PythonLexer()
formatter = HtmlFormatter(style="colorful")
html_code = highlight(code, lexer, formatter)

soup = BeautifulSoup(html_code, "html.parser")
pre = soup.find("pre")
lines = pre.decode_contents().splitlines()

for line in lines:
    line_soup = BeautifulSoup(line, "html.parser")
    clean_line = line_soup.prettify()
    print(clean_line)
