from typing import List, Tuple, Union
import numpy as np
from PIL import Image, ImageDraw
from pygments import lex
from pygments.formatters import ImageFormatter
from pygments.lexers import get_lexer_by_name
from enjoyn import GifAnimator, Preprocessor
def _custom_format(
formatter: ImageFormatter,
tokensource: List[Tuple],
max_line_length: int = None,
max_line_number: int = None,
) -> Image:
formatter._create_drawables(tokensource)
formatter._draw_line_numbers()
max_line_length = max_line_length or formatter.maxlinelength
max_line_number = max_line_number or formatter.maxlineno
image = Image.new(
"RGB",
formatter._get_image_size(max_line_length, max_line_number),
formatter.background_color,
)
formatter._paint_line_number_bg(image)
draw = ImageDraw.Draw(image)
# Highlight
if formatter.hl_lines:
x = (
formatter.image_pad
+ formatter.line_number_width
- formatter.line_number_pad
+ 1
)
recth = formatter._get_line_height()
rectw = image.size[0] - x
for linenumber in formatter.hl_lines:
y = formatter._get_line_y(linenumber - 1)
draw.rectangle([(x, y), (x + rectw, y + recth)], fill=formatter.hl_color)
for pos, value, font, text_fg, text_bg in formatter.drawables:
if text_bg:
text_size = draw.textsize(text=value, font=font)
draw.rectangle(
[pos[0], pos[1], pos[0] + text_size[0], pos[1] + text_size[1]],
fill=text_bg,
)
draw.text(pos, value, font=font, fill=text_fg)
return np.asarray(image)
def render_frame(
code: str,
formatter: ImageFormatter,
max_line_length: int = None,
max_line_number: int = None,
) -> Image:
lexer = get_lexer_by_name("python")
return _custom_format(
formatter,
lex(code, lexer),
max_line_length=max_line_length,
max_line_number=max_line_number,
)
code = """
from io import BytesIO
import matplotlib.pyplot as plt
import numpy as np
from enjoyn import GifAnimator, Preprocessor
def plot(i):
x = np.linspace(0, 2, 1000)
y = np.sin(2 * np.pi * (x - 0.01 * i))
plt.plot(x, y)
buf = BytesIO()
plt.savefig(buf)
plt.close()
return buf
output = GifAnimator(
preprocessor=plot,
items=np.arange(0, 200),
output_path="sine_wave.gif",
).compute()
""".strip()
formatter = ImageFormatter(
image_format="gif",
line_pad=8,
font_name=None,
line_number_bg=None,
line_number_fg=None,
encoding="utf-8",
)
longest_line = max(code.splitlines(), key=len) + " " * 12
max_line_length, _ = formatter.fonts.get_text_size(longest_line)
max_line_number = code.count("\n") + 1
items = [code[:i] for i in range(0, len(code) + 3, 3)]
items += [items[-1]] * 20 # pause at the end
preprocessor = Preprocessor(
func=render_frame,
args=(formatter,),
kwds=dict(max_line_length=max_line_length, max_line_number=max_line_number),
)
output = GifAnimator(
preprocessor=preprocessor, items=items, output_path="assets/pygments_code.gif"
).compute()
