atol bug-fix for leading underscores

- Support leading underscores (bug fix) for prefixed literal - Added relevant tests Signed-off-by: Joshua James Venter <[email protected]>
modular · Jul 9, 2024 · 6a0b673 · 6a0b673
1 parent d836be2
commit 6a0b673
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 9 deletions.
diff --git a/docs/changelog.md b/docs/changelog.md
@@ -415,6 +415,10 @@ what we publish.
 
 - `memcmp` no longer takes in `LegacyPointer`, instead, use `UnsafePointer`.
 
+- The `atol` function now correctly supports leading underscores,
+  (e.g.`atol("0x_ff", 0)`), when the appropriate base is specified or inferred
+  (base 0). non-base-10 integer literals as per Python's [Integer Literals](https://docs.python.org/3/reference/lexical_analysis.html#integers).
+
 ### ❌ Removed
 
 - It is no longer possible to cast (implicitly or explicitly) from `Reference`

diff --git a/stdlib/src/builtin/string.mojo b/stdlib/src/builtin/string.mojo
@@ -284,10 +284,7 @@ fn _atol(str_ref: StringRef, base: Int = 10) raises -> Int:
 
     var found_valid_chars_after_start = False
     var has_space_after_number = False
-    # single underscores are only allowed between digits
-    # starting "was_last_digit_undescore" to true such that
-    # if the first digit is an undesrcore an error is raised
-    var was_last_digit_undescore = True
+    var was_last_digit_undescore = real_base == 10
     for pos in range(start, str_len):
         var ord_current = int(buff[pos])
         if ord_current == ord_underscore:

diff --git a/stdlib/test/builtin/test_string.mojo b/stdlib/test/builtin/test_string.mojo
@@ -365,6 +365,9 @@ def test_atol():
     assert_equal(10, atol("0o12", 8))
     assert_equal(10, atol("0O12", 8))
     assert_equal(35, atol("Z", 36))
+    assert_equal(255, atol("0x_00_ff", 16))
+    assert_equal(18, atol("0b0001_0010", 2))
+    assert_equal(18, atol("0b_000_1001_0", 2))
 
     # Negative cases
     with assert_raises(
@@ -433,6 +436,14 @@ def test_atol_base_0():
 
     assert_equal(0, atol("0X0", base=0))
 
+    assert_equal(255, atol("0x_00_ff", base=0))
+
+    assert_equal(18, atol("0b_0001_0010", base=0))
+    assert_equal(18, atol("0b000_1001_0", base=0))
+
+    assert_equal(10, atol("0o_000_12", base=0))
+    assert_equal(10, atol("0o00_12", base=0))
+
     with assert_raises(
         contains="String is not convertible to integer with base 0: '  0x'"
     ):
@@ -453,11 +464,6 @@ def test_atol_base_0():
     ):
         _ = atol("0r100", base=0)
 
-    with assert_raises(
-        contains="String is not convertible to integer with base 0: '0b_0'"
-    ):
-        _ = atol("0b_0", base=0)
-
     with assert_raises(
         contains="String is not convertible to integer with base 0: '0xf__f'"
     ):