Enhance _parse_simple_number() to handle more cases

Author: vstinner

Lib/test/test_tomllib/test_misc.py

@@ -132,11 +132,26 @@ def test_parse_simple_number(self):
         parse_simple_number = tomllib._parser._parse_simple_number
         self.assertEqual(parse_simple_number("123", 0), (3, 123))
         self.assertEqual(parse_simple_number("123\n", 0), (3, 123))
+        self.assertEqual(parse_simple_number("123 456", 0), (3, 123))
+        self.assertEqual(parse_simple_number("+123\n", 0), (4, 123))
+        self.assertEqual(parse_simple_number("-123\n", 0), (4, -123))
         self.assertEqual(parse_simple_number("0\n", 0), (1, 0))
-
+        self.assertEqual(parse_simple_number("+0\n", 0), (2, 0))
+        self.assertEqual(parse_simple_number("-0\n", 0), (2, 0))
+        self.assertEqual(parse_simple_number("[23]\n", 1), (3, 23))
+        self.assertEqual(parse_simple_number("[23, 24]\n", 1), (3, 23))
+        self.assertEqual(parse_simple_number("[23]\n", 1), (3, 23))
+        self.assertEqual(parse_simple_number("{x = 42}\n", 5), (7, 42))
+
+        self.assertIsNone(parse_simple_number("+", 0), None)
+        self.assertIsNone(parse_simple_number("-", 0), None)
+        self.assertIsNone(parse_simple_number("+\n", 0), None)
+        self.assertIsNone(parse_simple_number("-\n", 0), None)
+        self.assertIsNone(parse_simple_number("+inf\n", 0), None)
+        self.assertIsNone(parse_simple_number("-nan\n", 0), None)
         self.assertIsNone(parse_simple_number("0123\n", 0))
-        self.assertIsNone(parse_simple_number("123-456\n", 0))
-        self.assertIsNone(parse_simple_number("123:456\n", 0))
+        self.assertIsNone(parse_simple_number("1979-05-27\n", 0))
+        self.assertIsNone(parse_simple_number("12:32:00\n", 0))
         self.assertIsNone(parse_simple_number("1.0\n", 0))
         self.assertIsNone(parse_simple_number("1_000\n", 0))
         self.assertIsNone(parse_simple_number("x123\n", 0))

Lib/tomllib/_parser.py

@@ -44,6 +44,8 @@
 KEY_INITIAL_CHARS: Final = BARE_KEY_CHARS | frozenset("\"'")
 HEXDIGIT_CHARS: Final = frozenset("abcdef" "ABCDEF" "0123456789")
 _DECDIGIT_CHARS: Final = frozenset("0123456789")
+_NUMBER_INITIAL_CHARS: Final = _DECDIGIT_CHARS | frozenset("+-")
+_NUMBER_END_CHARS: Final = frozenset(",]}") | TOML_WS_AND_NEWLINE
 
 BASIC_STR_ESCAPE_REPLACEMENTS: Final = MappingProxyType(
     {
@@ -672,19 +674,29 @@ def _parse_simple_number(
     src: str, pos: Pos
 ) -> None | tuple[Pos, int]:
     start = pos
-    src = src.rstrip()
     end = len(src)
+    end_chars = _NUMBER_END_CHARS
+    if src[pos] in '+-':
+        pos += 1
+        if pos >= end:
+            return None
+        if src[pos] not in _DECDIGIT_CHARS:
+            return None
+
+    if src[pos] == '0':
+        pos += 1
+        if pos < end and src[pos] not in end_chars:
+            return None
+        return pos, 0
+
     while src[pos] in _DECDIGIT_CHARS:
         pos += 1
         if pos >= end:
             break
     else:
-        if src[pos] != "\n":
+        if src[pos] not in end_chars:
             return None
-    digits = src[start:pos]
-    if digits.startswith("0") and len(digits) > 1:
-        return None
-    return pos, int(digits)
+    return pos, int(src[start:pos])
 
 
 def parse_value(
@@ -725,8 +737,9 @@ def parse_value(
     if char == "{":
         return parse_inline_table(src, pos, parse_float)
 
-    # Simple number parser avoiding regex
-    if char in _DECDIGIT_CHARS:
+    # First try a simple number parser which defers import tomllib._re
+    # to speed up tomllib import time
+    if char in _NUMBER_INITIAL_CHARS:
         res = _parse_simple_number(src, pos)
         if res is not None:
             return res