add custom turtle lib

myturtle.py

@@ -0,0 +1,139 @@
+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")
+