[mypyc] feat: new primitives for bytes.rjust and bytes.ljust

Author: BobTheBuidler

mypyc/lib-rt/bytes_ops.c

@@ -4,6 +4,7 @@
 
 #include <Python.h>
 #include "CPy.h"
+#include <string.h>
 
 // Returns -1 on error, 0 on inequality, 1 on equality.
 //
@@ -162,3 +163,51 @@ CPyTagged CPyBytes_Ord(PyObject *obj) {
     PyErr_SetString(PyExc_TypeError, "ord() expects a character");
     return CPY_INT_TAG;
 }
+
+PyObject *CPyBytes_Rjust(PyObject *self, Py_ssize_t width, PyObject *fillbyte) {
+    if (!PyBytes_Check(self)) {
+        PyErr_SetString(PyExc_TypeError, "self must be bytes");
+        return NULL;
+    }
+    if (!PyBytes_Check(fillbyte) || PyBytes_Size(fillbyte) != 1) {
+        PyErr_SetString(PyExc_TypeError, "fillbyte must be a single byte");
+        return NULL;
+    }
+    Py_ssize_t len = PyBytes_Size(self);
+    if (width <= len) {
+        return PyBytes_FromStringAndSize(PyBytes_AsString(self), len);
+    }
+    char fill = PyBytes_AsString(fillbyte)[0];
+    Py_ssize_t pad = width - len;
+    PyObject *result = PyBytes_FromStringAndSize(NULL, width);
+    if (!result) return NULL;
+    char *res_buf = PyBytes_AsString(result);
+    memset(res_buf, fill, pad);
+    memcpy(res_buf + pad, PyBytes_AsString(self), len);
+    return result;
+}
+
+PyObject *CPyBytes_Ljust(PyObject *self, Py_ssize_t width, PyObject *fillbyte) {
+    if (!PyBytes_Check(self)) {
+        PyErr_SetString(PyExc_TypeError, "self must be bytes");
+        return NULL;
+    }
+    if (!PyBytes_Check(fillbyte) || PyBytes_Size(fillbyte) != 1) {
+        PyErr_SetString(PyExc_TypeError, "fillbyte must be a single byte");
+        return NULL;
+    }
+    Py_ssize_t len = PyBytes_Size(self);
+    if (width <= len) {
+        return PyBytes_FromStringAndSize(PyBytes_AsString(self), len);
+    }
+    char fill = PyBytes_AsString(fillbyte)[0];
+    Py_ssize_t pad = width - len;
+    PyObject *result = PyBytes_FromStringAndSize(NULL, width);
+    if (!result) return NULL;
+    char *res_buf = PyBytes_AsString(result);
+    memcpy(res_buf, PyBytes_AsString(self), len);
+    memset(res_buf + len, fill, pad);
+    return result;
+}
+
+// ... existing code ...

mypyc/primitives/bytes_ops.py

@@ -126,3 +126,21 @@
     c_function_name="CPyBytes_Ord",
     error_kind=ERR_MAGIC,
 )
+
+# bytes.rjust(width, fillbyte=b' ')
+method_op(
+    name="rjust",
+    arg_types=[bytes_rprimitive, int_rprimitive, bytes_rprimitive],
+    return_type=bytes_rprimitive,
+    c_function_name="CPyBytes_Rjust",
+    error_kind=ERR_MAGIC,
+)
+
+# bytes.ljust(width, fillbyte=b' ')
+method_op(
+    name="ljust",
+    arg_types=[bytes_rprimitive, int_rprimitive, bytes_rprimitive],
+    return_type=bytes_rprimitive,
+    c_function_name="CPyBytes_Ljust",
+    error_kind=ERR_MAGIC,
+)

mypyc/test-data/irbuild-bytes.test

@@ -185,3 +185,69 @@ L0:
     r10 = CPyBytes_Build(2, var, r9)
     b4 = r10
     return 1
+
+[case testBytesRjustDefault]
+def f(b: bytes) -> bytes:
+    return b.rjust(6)
+[out]
+def f(b):
+    b :: bytes
+    r0 :: bytes
+L0:
+    r0 = b.rjust(6, b' ')
+    return r0
+
+[case testBytesRjustCustom]
+def f(b: bytes) -> bytes:
+    return b.rjust(8, b'0')
+[out]
+def f(b):
+    b :: bytes
+    r0 :: bytes
+L0:
+    r0 = b.rjust(8, b'0')
+    return r0
+
+[case testBytesLjustDefault]
+def f(b: bytes) -> bytes:
+    return b.ljust(7)
+[out]
+def f(b):
+    b :: bytes
+    r0 :: bytes
+L0:
+    r0 = b.ljust(7, b' ')
+    return r0
+
+[case testBytesLjustCustom]
+def f(b: bytes) -> bytes:
+    return b.ljust(10, b'_')
+[out]
+def f(b):
+    b :: bytes
+    r0 :: bytes
+L0:
+    r0 = b.ljust(10, b'_')
+    return r0
+
+[case testBytesRjustNoPad]
+def f(b: bytes) -> bytes:
+    return b.rjust(2)
+[out]
+def f(b):
+    b :: bytes
+    r0 :: bytes
+L0:
+    r0 = b.rjust(2, b' ')
+    return r0
+
+[case testBytesLjustNoPad]
+def f(b: bytes) -> bytes:
+    return b.ljust(1)
+[out]
+def f(b):
+    b :: bytes
+    r0 :: bytes
+L0:
+    r0 = b.ljust(1, b' ')
+    return r0

mypyc/test-data/run-bytes.test

@@ -374,3 +374,41 @@ class subbytearray(bytearray):
 [file userdefinedbytes.py]
 class bytes:
     pass
+
+[case testBytesRjustLjust]
+def rjust_bytes(b: bytes, width: int, fill: bytes = b' ') -> bytes:
+    return b.rjust(width, fill)
+
+def ljust_bytes(b: bytes, width: int, fill: bytes = b' ') -> bytes:
+    return b.ljust(width, fill)
+
+def test_rjust_with_default_fill() -> None:
+    assert rjust_bytes(b'abc', 6) == b'   abc'
+    assert rjust_bytes(b'abc', 3) == b'abc'
+    assert rjust_bytes(b'abc', 2) == b'abc'
+    assert rjust_bytes(b'', 4) == b'    '
+
+def test_rjust_with_custom_fill() -> None:
+    assert rjust_bytes(b'abc', 6, b'0') == b'000abc'
+    assert rjust_bytes(b'abc', 5, b'_') == b'__abc'
+    assert rjust_bytes(b'abc', 3, b'X') == b'abc'
+
+def test_ljust_with_default_fill() -> None:
+    assert ljust_bytes(b'abc', 6) == b'abc   '
+    assert ljust_bytes(b'abc', 3) == b'abc'
+    assert ljust_bytes(b'abc', 2) == b'abc'
+    assert ljust_bytes(b'', 4) == b'    '
+
+def test_ljust_with_custom_fill() -> None:
+    assert ljust_bytes(b'abc', 6, b'0') == b'abc000'
+    assert ljust_bytes(b'abc', 5, b'_') == b'abc__'
+    assert ljust_bytes(b'abc', 3, b'X') == b'abc'
+
+def test_edge_cases() -> None:
+    assert rjust_bytes(b'abc', 0) == b'abc'
+    assert ljust_bytes(b'abc', 0) == b'abc'
+    # fillbyte must be length 1
+    with assertRaises(TypeError):
+        rjust_bytes(b'abc', 5, b'')
+    with assertRaises(TypeError):
+        ljust_bytes(b'abc', 5, b'12')