[python] Return value for turtle.pensize()

apps/code/script_template.cpp

@@ -73,6 +73,8 @@ def spiral(N_iteration):
 
 constexpr ScriptTemplate kochScriptTemplate("koch.py", "\x01" R"(from turtle import *
 def koch(N_iteration):
+# Save pencil size and set it to 1
+  prevSize = pensize()
   pensize(1)
 # Define inner function for recursion
   def _koch(n, l):
@@ -87,7 +89,10 @@ def koch(N_iteration):
       left(60)
       _koch(n-1,l/3)
 # Call inner function
-  _koch(N_iteration,140))");
+  _koch(N_iteration,140)
+# Restore pencil size
+  pensize(prevSize))");
+
 const ScriptTemplate * ScriptTemplate::Empty() {
   return &emptyScriptTemplate;
 }

python/port/modturtle.c

@@ -12,7 +12,7 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_1(turtle_speed_obj, turtle_speed);
 
 STATIC MP_DEFINE_CONST_FUN_OBJ_0(turtle_pendown_obj, turtle_pendown);
 STATIC MP_DEFINE_CONST_FUN_OBJ_0(turtle_penup_obj, turtle_penup);
-STATIC MP_DEFINE_CONST_FUN_OBJ_1(turtle_pensize_obj, turtle_pensize);
+STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(turtle_pensize_obj, 0, 1, turtle_pensize);
 
 STATIC MP_DEFINE_CONST_FUN_OBJ_3(turtle_color_obj, turtle_color);
 

python/port/modturtle.h

@@ -10,7 +10,7 @@ mp_obj_t turtle_speed(mp_obj_t speed);
 
 mp_obj_t turtle_pendown();
 mp_obj_t turtle_penup();
-mp_obj_t turtle_pensize(mp_obj_t size);
+mp_obj_t turtle_pensize(size_t n_args, const mp_obj_t *args);
 
 mp_obj_t turtle_color(mp_obj_t r, mp_obj_t g, mp_obj_t b);
 
@@ -18,4 +18,6 @@ mp_obj_t turtle_showturtle();
 mp_obj_t turtle_hideturtle();
 
 mp_obj_t turtle___init__();
+
+void turtle_initpen(int size);
 void turtle_deinit();

python/port/modturtle_impl.cpp

@@ -171,35 +171,20 @@ mp_obj_t turtle_penup() {
   return mp_const_none;
 }
 
-mp_obj_t turtle_pensize(mp_obj_t size) {
-  int s = mp_obj_get_int(size);
-  if (s < 1) {
-    s = 1;
-  }
-  else if (s > 10) {
-    s = 10;
+mp_obj_t turtle_pensize(size_t n_args, const mp_obj_t *args) {
+  if (n_args == 0) {
+    return MP_OBJ_NEW_SMALL_INT(t_dotsize);
   }
 
-  if (t_dot) {
-    delete[] t_dot;
+  int size = mp_obj_get_int(args[0]);
+  if (size < 1) {
+    size = 1;
+  }
+  else if (size > 10) {
+    size = 10;
   }
-  t_dot = new uint8_t[s*s];
-  t_dotsize = s;
 
-  float middle = s / 2;
-  for (int j = 0; j < s; j++) {
-    for (int i = 0; i < s; i++) {
-      float distance = sqrt((j - middle)*(j - middle) + (i - middle)*(i - middle)) / (middle+1);
-      int value = distance * distance * 255;
-      if (value < 0) {
-        value = 0;
-      }
-      else if (value > 255) {
-        value = 255;
-      }
-      t_dot[j*s + i] = value;
-    }
-  }
+  turtle_initpen(size);
 
   return mp_const_none;
 }
@@ -244,11 +229,34 @@ mp_obj_t turtle___init__() {
   t_mileage = 0;
   t_hidden = false;
 
-  turtle_pensize(MP_OBJ_NEW_SMALL_INT(5));
+  turtle_initpen(5);
 
   return mp_const_none;
 }
 
+void turtle_initpen(int size) {
+  if (t_dot) {
+    delete[] t_dot;
+  }
+  t_dot = new uint8_t[size*size];
+  t_dotsize = size;
+
+  float middle = size / 2;
+  for (int j = 0; j < size; j++) {
+    for (int i = 0; i < size; i++) {
+      float distance = sqrt((j - middle)*(j - middle) + (i - middle)*(i - middle)) / (middle+1);
+      int value = distance * distance * 255;
+      if (value < 0) {
+        value = 0;
+      }
+      else if (value > 255) {
+        value = 255;
+      }
+      t_dot[j*size + i] = value;
+    }
+  }
+}
+
 void turtle_deinit() {
   delete[] t_underneath;
   t_underneath = nullptr;