gh-101178: refactor base64.b85encode to be memory friendly

Lib/base64.py

@@ -298,27 +298,12 @@ def b16decode(s, casefold=False):
 
 def _85encode(b, chars, chars2, pad=False, foldnuls=False, foldspaces=False):
  # Helper function for a85encode and b85encode
+ # chars2 is now unused
  if not isinstance(b, bytes_types):
  b = memoryview(b).tobytes()
 
- padding = (-len(b)) % 4
- if padding:
- b = b + b'\0' * padding
- words = struct.Struct('!%dI' % (len(b) // 4)).unpack(b)
-
- chunks = [b'z' if foldnuls and not word else
- b'y' if foldspaces and word == 0x20202020 else
- (chars2[word // 614125] +
- chars2[word // 85 % 7225] +
- chars[word % 85])
- for word in words]
-
- if padding and not pad:
- if chunks[-1] == b'z':
- chunks[-1] = chars[0] * 5
- chunks[-1] = chunks[-1][:-padding]
-
- return b''.join(chunks)
+ return binascii._b2a_base85(b, chars=chars, pad=pad,
+ foldnuls=foldnuls, foldspaces=foldspaces)
 
 def a85encode(b, *, foldspaces=False, wrapcol=0, pad=False, adobe=False):
  """Encode bytes-like object b using Ascii85 and return a bytes object.
@@ -337,14 +322,13 @@ def a85encode(b, *, foldspaces=False, wrapcol=0, pad=False, adobe=False):
  adobe controls whether the encoded byte sequence is framed with <~ and ~>,
  which is used by the Adobe implementation.
  """
- global _a85chars, _a85chars2
+ global _a85chars
  # Delay the initialization of tables to not waste memory
  # if the function is never called
- if _a85chars2 is None:
+ if _a85chars is None:
  _a85chars = [bytes((i,)) for i in range(33, 118)]
- _a85chars2 = [(a + b) for a in _a85chars for b in _a85chars]
 
- result = _85encode(b, _a85chars, _a85chars2, pad, True, foldspaces)
+ result = _85encode(b, b''.join(_a85chars), None, pad, True, foldspaces)
 
  if adobe:
  result = _A85START + result
@@ -445,13 +429,12 @@ def b85encode(b, pad=False):
  If pad is true, the input is padded with b'\\0' so its length is a multiple of
  4 bytes before encoding.
  """
- global _b85chars, _b85chars2
+ global _b85chars
  # Delay the initialization of tables to not waste memory
  # if the function is never called
- if _b85chars2 is None:
+ if _b85chars is None:
  _b85chars = [bytes((i,)) for i in _b85alphabet]
- _b85chars2 = [(a + b) for a in _b85chars for b in _b85chars]
- return _85encode(b, _b85chars, _b85chars2, pad)
+ return _85encode(b, _b85alphabet, None, pad)
 
 def b85decode(b):
  """Decode the base85-encoded bytes-like object or ASCII string b

Lib/test/test_base64.py

@@ -528,6 +528,7 @@ def test_b85encode(self):
  b"""0123456789!@#0^&*();:<>,. []{}""":
  b"""VPa!sWoBn+X=-b1ZEkOHadLBXb#`}nd3r%YLqtVJM@UIZOH55pPf$@("""
  b"""Q&d$}S6EqEFflSSG&MFiI5{CeBQRbjDkv#CIy^osE+AW7dwl""",
+ b"paddu\xc7": b'aA9O*b;k',
  b'no padding..': b'Zf_uPVPs@!Zf7no',
  b'zero compression\x00\x00\x00\x00': b'dS!BNAY*TBaB^jHb7^mG00000',
  b'zero compression\x00\x00\x00': b'dS!BNAY*TBaB^jHb7^mG0000',

Modules/binascii.c

@@ -1239,13 +1239,108 @@ binascii_b2a_qp_impl(PyObject *module, Py_buffer *data, int quotetabs,
  return rv;
 }
 
+/*[clinic input]
+binascii._b2a_base85
+
+ data: Py_buffer
+ chars: Py_buffer
+ pad: bool = False
+ foldnuls: bool = False
+ foldspaces: bool = False
+
+Utility method used by the base64 module to encode a85/b85 data
+
+ data: bytes
+ chars: 85 bytes conversion table
+ pad: use NULL-paded input if necessary
+ foldnuls: replace NULL chunks by 'z'
+ foldspaces: replace space-only chucks by 'y'
+
+[clinic start generated code]*/
+
+static PyObject *
+binascii__b2a_base85_impl(PyObject *module, Py_buffer *data,
+ Py_buffer *chars, int pad, int foldnuls,
+ int foldspaces)
+/*[clinic end generated code: output=cefe84c300ad7314 input=3c8faf77b992dcc2]*/
+{
+ if (chars->len != 85) {
+ PyErr_SetString(PyExc_ValueError,
+ "chars must be exactly 85 bytes long");
+ return NULL;
+ }
+
+ _PyBytesWriter writer;
+ _PyBytesWriter_Init(&writer);
+
+ const size_t bin_len = data->len;
+
+ // Allocate up to maxium encoded length, adjusted at end
+ const size_t ascii_len = ((bin_len + 3) / 4) * 5;
+
+ unsigned char *ascii_data = _PyBytesWriter_Alloc(&writer, ascii_len);
+ if (ascii_data == NULL) {
+ PyErr_NoMemory();
+ return NULL;
+ }
+
+ const unsigned char *table = chars->buf;
+ const unsigned char *bin_data = data->buf;
+
+ size_t i = 0 ;
+ int padding = 0;
+
+ // Conversion largely inspired from git base85 implementation
+ while (i < bin_len) {
+ // Translate each 4 byte chunk to 32bit integer
+ uint32_t value = 0;
+ for (int cnt = 24; cnt >= 0; cnt -= 8) {
+ value |= bin_data[i] << cnt;
+ if (++i == bin_len) {
+ // Number of bytes under the 4 bytes rounded value
+ padding = cnt / 8;
+ break;
+ }
+ }
+
+ // Handle NULL only and space-only cases (specific to ASCII85)
+ if (foldnuls && value == 0) {
+ *ascii_data++ = 'z';
+ }
+ else if (foldspaces && value == 0x20202020) {
+ *ascii_data++ = 'y';
+ }
+ else {
+ for (int j = 4; j >= 0; j--) {
+ ascii_data[j] = table[value % 85];
+ value /= 85;
+ }
+ ascii_data += 5;
+ }
+ }
+
+ // Expand the last folded null in case it did not fill a full chunk
+ if (padding && !pad && foldnuls && ascii_data[-1] == 'z') {
+ ascii_data--;
+ memset(ascii_data, table[0], 5);
+ ascii_data += 5;
+ }
+
+ if (!pad) {
+ ascii_data -= padding;
+ }
+
+ return _PyBytesWriter_Finish(&writer, ascii_data);
+}
+
 /* List of functions defined in the module */
 
 static struct PyMethodDef binascii_module_methods[] = {
  BINASCII_A2B_UU_METHODDEF
  BINASCII_B2A_UU_METHODDEF
  BINASCII_A2B_BASE64_METHODDEF
  BINASCII_B2A_BASE64_METHODDEF
+ BINASCII__B2A_BASE85_METHODDEF
  BINASCII_A2B_HEX_METHODDEF
  BINASCII_B2A_HEX_METHODDEF
  BINASCII_HEXLIFY_METHODDEF