gh-158585: Optimize bytes object creation (#158664)

Add bytes_alloc() helper function. bytes_alloc() and
_PyBytes_FromStringAndSize() have less conditional branches than
previous code, and so should be a little bit faster.

* Rename existing bytes_alloc() to bytes_type_alloc().
* Replace _PyBytes_FromSize(calloc=1) with _PyBytes_FromSizeZero().
* Add set_ob_shash_unsafe(): similar to set_ob_shash() but don't use
  an atomic operation on Free Threading. Use this new function
  on newly allocated bytes objects and in bytes_resize_inplace().
* Replace PyBytes_FromStringAndSize(NULL, size) with
  bytes_alloc(size).
* Use bytes_alloc() in PyBytes_FromString() and _PyBytes_Repeat().
* Add _PyBytes_FromStringAndSize(): similar to
  PyBytes_FromStringAndSize(), but str must not be NULL.
* Replace PyBytes_FromStringAndSize() with
  _PyBytes_FromStringAndSize().
* Test that PyBytes_FromString() and PyBytes_FromStringAndSize()
  return singletons for 0 or 1 bytes.
* Fix test_tracemalloc.test_set_traceback_limit().

Author: vstinner

Lib/test/test_bytes.py

@@ -410,6 +410,8 @@ def test_concat(self):
         self.assertRaises(TypeError, lambda: b1 + "def")
         self.assertRaises(TypeError, lambda: "abc" + b2)
 
+        self.assertEqual(self.type2test() + memoryview(b''), b'')
+
     def test_repeat(self):
         for b in b"abc", self.type2test(b"abc"):
             self.assertEqual(b * 3, b"abcabcabc")
@@ -1447,6 +1449,45 @@ class SubBytes(bytes):
         self.assertNotEqual(id(s), id(1 * s))
         self.assertNotEqual(id(s), id(s * 2))
 
+    @support.cpython_only
+    def test_concat_cpython(self):
+        # Test optimizations
+        empty = b''
+        abc = b"abc"
+        self.assertIs(abc + empty, abc)
+        self.assertIs(empty + abc, abc)
+        self.assertIs(empty + bytearray(b''), empty)
+        self.assertIs(empty + memoryview(b''), empty)
+
+        class Subclass(bytes):
+            pass
+
+        empty_subclass = Subclass(b"")
+        self.assertIs(empty + empty_subclass, empty)
+        self.assertIs(empty_subclass + empty, empty)
+        self.assertIs(abc + empty_subclass, abc)
+        self.assertIs(empty_subclass + abc, abc)
+
+        def assert_copy(result, expected):
+            self.assertIsNot(result, expected)
+            self.assertEqual(type(result), bytes)
+            self.assertEqual(result, expected)
+
+        # Copy the string if it's a subclass
+        def_subclass = Subclass(b"def")
+        assert_copy(def_subclass + empty, b'def')
+        assert_copy(empty + def_subclass, b'def')
+        assert_copy(abc + def_subclass, b'abcdef')
+        assert_copy(def_subclass + abc, b'defabc')
+
+    @support.cpython_only
+    def test_repeat_cpython(self):
+        # Test optimizations
+        for b in b'', b'x', b"abc":
+            self.assertIs(b * 1, b)
+            self.assertIs(b * 0, b'')
+            self.assertIs(b * -23, b'')
+
 
 class ByteArrayTest(BaseBytesTest, unittest.TestCase):
     type2test = bytearray

Lib/test/test_capi/test_bytes.py

@@ -61,10 +61,12 @@ def test_fromstringandsize(self):
         self.assertEqual(fromstringandsize(b'abc'), b'abc')
         self.assertEqual(fromstringandsize(b'abc', 2), b'ab')
         self.assertEqual(fromstringandsize(b'abc\0def'), b'abc\0def')
-        self.assertEqual(fromstringandsize(b'a'), b'a')
-        self.assertEqual(fromstringandsize(b'a', 1), b'a')
-        self.assertEqual(fromstringandsize(b'', 0), b'')
-        self.assertEqual(fromstringandsize(NULL, 0), b'')
+
+        # Return singleton for 0 or 1 bytes
+        self.assertIs(fromstringandsize(b'', 0), b'')
+        self.assertIs(fromstringandsize(NULL, 0), b'')
+        self.assertIs(fromstringandsize(b'a', 1), b'a')
+
         self.assertEqual(len(fromstringandsize(NULL, 3)), 3)
         self.assertRaises((MemoryError, OverflowError),
                           fromstringandsize, NULL, PY_SSIZE_T_MAX)
@@ -79,7 +81,10 @@ def test_fromstring(self):
         fromstring = _testlimitedcapi.bytes_fromstring
 
         self.assertEqual(fromstring(b'abc\0def'), b'abc')
-        self.assertEqual(fromstring(b''), b'')
+
+        # Return singleton for 0 or 1 bytes
+        self.assertIs(fromstring(b''), b'')
+        self.assertIs(fromstring(b'a'), b'a')
 
         # CRASHES fromstring(NULL)
 

Lib/test/test_tracemalloc.py

@@ -41,6 +41,8 @@ def get_frames(nframe, lineno_delta):
 def allocate_bytes(size):
     nframe = tracemalloc.get_traceback_limit()
     bytes_len = (size - EMPTY_STRING_SIZE)
+    if bytes_len < 0:
+        raise ValueError(f"{size} bytes is too small")
     frames = get_frames(nframe, 1)
     data = b'x' * bytes_len
     return data, tracemalloc.Traceback(frames, min(len(frames), nframe))
@@ -158,7 +160,7 @@ def test_new_reference(self):
         self.assertEqual(traceback, obj_traceback)
 
     def test_set_traceback_limit(self):
-        obj_size = 10
+        obj_size = 123
 
         tracemalloc.stop()
         self.assertRaises(ValueError, tracemalloc.start, -1)