[파이썬] 문제 : 터틀(turtle) 도형의 종류와 가로,세로를 입력 받고 그리기

import turtle
from time import*
import datetime

screen = turtle.Screen()
screen.setup(500, 300)        # 화면 크기 

t = turtle.Turtle()
t.shape('turtle')
# t.speed(10)                  # 그리는 속도. 빠르게

def draw(n, w,  h):           # 그리기 함수
    for i in range(n):        # n 각형. 수 만큼 반복
        if i%2 == 0:          # 0,2,4와 같이 짝수 일때. 가로일때
            t.forward(w)       # w 만큼 그리기
        else:                    # 홀수면, 세로일때
            t.forward(h)          # h 만큼 그리기
        if n ==  3 and i < n-1:  # 마지막에는 거북이의 각도를 바꾸지 않음. 삼각형일때 각도를 left 로 바꿈
            t.left(360/n)         # 360도를 n으로 나누면 한번 바꾸는 각도가 됨
        elif n == 4 and i < n-1: # 마지막에는 거북이의 각도를 바꾸지 않음. 사각형일때 각도를 right로 바꿈
            t.right(360/n)


s = input('도형의 종류 : ')            # 도형 종료 입력 받음

if s == '삼각형':                      # 삼각형일때 
    w = int(input('한 변의 길이 : '))  # 한변의 길이 입력 받음
    h = w                              # h도 w로 같이 함 
    n = 3                              # 함수에 인수로 넘겨줄 도형 종료
elif s == '사각형':                   # 사각형일때 
    w = int(input('가로 길이 : '))     # w, h를 각각 입력 받음
    h = int(input('세로 길이 : '))
    n = 4
    
draw(n, w, h)          # 그리기 함수에 도형 종료, 가로와 세로를 인수로 넘겨 줌
    
turtle.done()