xu3kev's picture
add custom turtle lib
be3d62e
raw
history blame
No virus
3.45 kB
import re
from PIL import Image
from matplotlib import pyplot as plt
import matplotlib as mpl
import matplotlib.style as mplstyle
import numpy as np
mpl.use('Agg')
mpl.rcParams["path.simplify_threshold"] = 0.0
mpl.rcParams['agg.path.chunksize'] = 10000
mplstyle.use('fast')
HALF_INF = 63
INF = 126
EPS_DIST = 1/20
EPS_ANGLE = 2.86
class Turtle:
def __init__(self, ax=None):
self.x = 0
self.y = 0
self.heading = 0
if ax is None:
self.fig, self.ax = plt.subplots(1, 1, figsize=(20, 20))
else:
self.ax = ax
self.ax.set_xlim(-50, 50)
self.ax.set_ylim(-50, 50)
self.ax.set_xticklabels([])
self.ax.set_yticklabels([])
self.ax.set_xticks([])
self.ax.set_yticks([])
self.ax.spines['top'].set_visible(False)
self.ax.spines['bottom'].set_visible(False)
self.ax.spines['left'].set_visible(False)
self.ax.spines['right'].set_visible(False)
self.is_down = True
def forward(self, dist):
x0, y0 = self.x, self.y
x1 = x0 + dist * np.cos(self.heading)
y1 = y0 + dist * np.sin(self.heading)
if self.is_down:
self.ax.plot([x0, x1], [y0, y1], color='black', linewidth=3)
self.x = x1
self.y = y1
def left(self, angle):
self.heading += angle * np.pi / 180
def right(self, angle):
self.heading -= angle * np.pi / 180
def penup(self):
self.is_down = False
def pendown(self):
self.is_down = True
def save(self, path):
self.fig.canvas.draw()
pil_img = Image.frombytes('RGB', self.fig.canvas.get_width_height(), self.fig.canvas.tostring_rgb())
if path:
pil_img.save(path)
return pil_img
class _TurtleState:
def __init__(self, turtle):
self.turtle = turtle
self.position = None
self.heading = None
self.pen_status = None
def __enter__(self):
self.position = (self.turtle.x, self.turtle.y)
self.heading = self.turtle.heading
self.pen_status = self.turtle.is_down
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.turtle.penup()
self.turtle.x, self.turtle.y = self.position
self.turtle.heading = self.heading
if self.pen_status:
self.turtle.pendown()
if __name__ == "__main__":
turtle = Turtle()
def forward(dist):
turtle.forward(dist)
def left(angle):
turtle.left(angle)
def right(angle):
turtle.right(angle)
def penup():
turtle.penup()
def pendown():
turtle.pendown()
def save(path):
turtle.save(path)
def fork_state():
"""
Clone the current state of the turtle.
Usage:
with clone_state():
forward(100)
left(90)
forward(100)
"""
return turtle._TurtleState(turtle)
# Example usage
def example_plot():
forward(5)
with fork_state():
forward(10)
left(90)
forward(10)
with fork_state():
right(90)
forward(20)
left(90)
forward(10)
left(90)
forward(10)
right(90)
forward(50)
save("test.png")