diff --git a/setup.py b/setup.py index 822835c..4e5f9b4 100644 --- a/setup.py +++ b/setup.py @@ -7,7 +7,7 @@ setuptools.setup( name="spyrograph", - version="0.21.0", + version="0.22.0", author="Chris Greening", author_email="chris@christophergreening.com", description="Library for drawing spirographs in Python", diff --git a/spyrograph/core/_misc.py b/spyrograph/core/_misc.py index 2a0a5b6..df15091 100644 --- a/spyrograph/core/_misc.py +++ b/spyrograph/core/_misc.py @@ -9,6 +9,11 @@ import numpy as np +try: + from PIL import ImageGrab +except ImportError: + ImageGrab = None + def _validate_theta( thetas: List[Number], theta_start: Number, theta_stop: Number, theta_step: Number @@ -50,6 +55,28 @@ def _set_int_to_list(input_val: Union[Number, List[Number]]) -> List[Number]: input_val = [input_val] return input_val +def _save_trace(screen: "turtle.Turtle", fpath: str): + """Save trace to PNG using PIL""" + # pylint: disable=invalid-name + if ImageGrab is None: + raise ImportError(( + "PIL is required but is not installed on your machine, " + "please install and try again" + )) + canvas = screen.getcanvas() + root = canvas.winfo_toplevel() + root.update() + x0 = root.winfo_rootx() + y0 = root.winfo_rooty() + time.sleep(1) + image = ImageGrab.grab(( + x0+8, + y0+8, + x0 + root.winfo_width()-8, + y0 + root.winfo_height()-8 + )) + image.save(fpath) + def _draw_animation( shapes_arr, screen_size: Tuple[Number, Number] = (1000, 1000), screen_color: str = "white", exit_on_click: bool = False, diff --git a/spyrograph/core/_trochoid.py b/spyrograph/core/_trochoid.py index 6b65d10..5ebd83c 100644 --- a/spyrograph/core/_trochoid.py +++ b/spyrograph/core/_trochoid.py @@ -14,7 +14,8 @@ import numpy as np from spyrograph.core._misc import ( - _get_products_of_inputs, _validate_only_one_iterable, _draw_animation, _validate_theta + _get_products_of_inputs, _validate_only_one_iterable, _draw_animation, + _validate_theta, _save_trace ) try: @@ -82,6 +83,11 @@ def __init__( self.y = np.array([self._calculate_y(theta) for theta in self.thetas]) self.x += self.origin[0] self.y += self.origin[1] + self.min_x = min(self.x) + self.max_x = max(self.x) + self.min_y = min(self.y) + self.max_y = max(self.y) + self.coords = list(zip(self.x, self.y, self.thetas)) def transform(self, x: Number = 0, y: Number = 0) -> "_Trochoid": @@ -212,14 +218,52 @@ def plot(self, **kwargs) -> Tuple["matplotlib.matplotlib.Figure", "matplotlib.ax plt.show() return fig, ax + def save_png( + self, fpath: str, screen_size: Tuple[Number, Number] = None, + screen_color: str = "white", color: str = "black", width: Number = 1, + screen: "turtle.Screen" = None, screen_coords = (0, 0), padding = 100 + ) -> None: + """ + Save the shape as a PNG file. + + Parameters + ---------- + fpath : str + The file path where the PNG file will be saved. + screen_size : Tuple[Number, Number], optional + The width and height of the turtle screen. Default is None. + screen_color : str, optional + The background color of the turtle screen. Default is "white". + color : str, optional + The color of the shape. Default is "black". + width : Number, optional + The width of the shape lines. Default is 1. + screen : "turtle.Screen", optional + The turtle screen object to draw the shape on. Default is None. + screen_coords : Tuple[Number, Number], optional + The x and y coordinates of the top-left corner of the turtle screen. Default is (0, 0). + padding : int, optional + The padding around the shape in the final PNG image. Default is 100. + + Examples + -------- + >>> shape = Trochoid(R=250, r=179, d=233, thetas=np.arange(0, 60, .01)) + >>> shape.save_png("spirograph.png", width=2) + """ + screen, _ = self.trace( + screen_size=screen_size, screen_color=screen_color, color=color, + width=width, screen=screen, screen_coords=screen_coords, padding=padding + ) + _save_trace(screen, fpath) + def trace( - self, screen_size: Tuple[Number, Number] = (1000, 1000), + self, screen_size: Tuple[Number, Number] = None, screen_color: str = "white", exit_on_click: bool = False, color: str = "black", width: Number = 1, hide_turtle: bool = True, show_circles: bool = False, frame_pause: Number = 0, screen: "turtle.Screen" = None, circle_color: str = "black", show_full_path: bool = False, full_path_color: str = "grey", - repeat: bool = False, screen_coords = (0, 0) + repeat: bool = False, screen_coords = (0, 0), padding: Number = 100 ) -> "turtle.Screen": """ Trace the shape using the turtle graphics library and return the turtle.Screen object. @@ -273,7 +317,7 @@ def trace( >>> screen = shape.trace(show_circles=True, exit_on_click=True) """ # pylint: disable=no-member,too-many-locals - screen = self._init_screen(screen, screen_size, screen_color, screen_coords) + screen = self._init_screen(screen, screen_size, screen_color, screen_coords, padding) turtle.tracer(False) turtles = self._init_turtles(color, circle_color, full_path_color, hide_turtle, width) @@ -308,7 +352,7 @@ def trace( turtle.update() if not repeat: break - turtles.shape_turtle.clear() + # turtles.shape_turtle.clear() if exit_on_click: turtle.exitonclick() return screen, turtles @@ -526,11 +570,16 @@ def _show_full_path(self, pre_draw_turtle: "turtle.Turtle") -> None: def _init_screen( self, screen: "turtle.Screen", screen_size: Tuple[Number, Number], - screen_color: str, screen_coords: Tuple[Number, Number] + screen_color: str, screen_coords: Tuple[Number, Number], padding: Number ) -> "turtle.Screen": """Initializes the turtle screen with the given size and color""" if screen is None: screen = turtle.Screen() + if screen_size is None: + screen_size = ( + self.max_x - self.min_x + padding, + self.max_y - self.min_y + padding + ) screen.setup(*screen_size) screen.bgcolor(screen_color) canvas = screen.getcanvas()