From b1115a4c8fb53280c9728af9c298077c1b67a405 Mon Sep 17 00:00:00 2001
From: albertexye <111392956+albertexye@users.noreply.github.com>
Date: Mon, 1 Apr 2024 14:41:45 -0400
Subject: [PATCH] Check the size of fillchar passed to str methods (ljust,
 rjust) (#236)

* check the size of fillchar passed to str methods (ljust, rjust)

* count characters using u8_length instead of size
---
 src/pocketpy.cpp |  2 ++
 tests/04_str.py  | 10 ++++++++++
 2 files changed, 12 insertions(+)

diff --git a/src/pocketpy.cpp b/src/pocketpy.cpp
index 37284b63..27455846 100644
--- a/src/pocketpy.cpp
+++ b/src/pocketpy.cpp
@@ -734,6 +734,7 @@ void init_builtins(VM* _vm) {
         int delta = width - self.u8_length();
         if(delta <= 0) return args[0];
         const Str& fillchar = CAST(Str&, args[2]);
+        if (fillchar.u8_length() != 1) vm->TypeError("The fill character must be exactly one character long");
         SStream ss;
         ss << self;
         for(int i=0; i<delta; i++) ss << fillchar;
@@ -747,6 +748,7 @@ void init_builtins(VM* _vm) {
         int delta = width - self.u8_length();
         if(delta <= 0) return args[0];
         const Str& fillchar = CAST(Str&, args[2]);
+        if (fillchar.u8_length() != 1) vm->TypeError("The fill character must be exactly one character long");
         SStream ss;
         for(int i=0; i<delta; i++) ss << fillchar;
         ss << self;
diff --git a/tests/04_str.py b/tests/04_str.py
index 4c3d5ef8..3a588229 100644
--- a/tests/04_str.py
+++ b/tests/04_str.py
@@ -157,8 +157,18 @@ assert b[5:2:-2] == [',', 'l']
 a = '123'
 assert a.rjust(5) == '  123'
 assert a.rjust(5, '0') == '00123'
+try:
+    a.rjust(5, '00')
+    exit(1)
+except TypeError:
+    pass
 assert a.ljust(5) == '123  '
 assert a.ljust(5, '0') == '12300'
+try:
+    a.ljust(5, '00')
+    exit(1)
+except TypeError:
+    pass
 
 assert '\x30\x31\x32' == '012'