Merge pull request #11 from ngoldbaum/add-asciidtype

Fix bugs around null termination bytes

Author: ngoldbaum

asciidtype/asciidtype/src/casts.c

@@ -25,7 +25,7 @@ ascii_to_ascii_resolve_descriptors(PyObject *NPY_UNUSED(self),
     }
     else {
         Py_INCREF(given_descrs[1]);
-        loop_descrs[1] = given_descrs[0];
+        loop_descrs[1] = given_descrs[1];
     }
 
     if (((ASCIIDTypeObject *)loop_descrs[0])->size ==
@@ -38,35 +38,9 @@ ascii_to_ascii_resolve_descriptors(PyObject *NPY_UNUSED(self),
 }
 
 static int
-ascii_to_ascii_contiguous(PyArrayMethod_Context *context, char *const data[],
-                          npy_intp const dimensions[],
-                          npy_intp const NPY_UNUSED(strides[]),
-                          NpyAuxData *NPY_UNUSED(auxdata))
-{
-    PyArray_Descr **descrs = context->descriptors;
-    // for contiguous assignment the sizes of the two dtypes should be
-    // the same, consider adding an assert to check?
-    long size = ((ASCIIDTypeObject *)descrs[0])->size;
-
-    npy_intp N = dimensions[0] * size;
-    char *in = data[0];
-    char *out = data[1];
-
-    while (N--) {
-        *out = *in;
-        out++;
-        in++;
-    }
-
-    return 0;
-}
-
-static int
-ascii_to_ascii_strided_or_unaligned(PyArrayMethod_Context *context,
-                                    char *const data[],
-                                    npy_intp const dimensions[],
-                                    npy_intp const strides[],
-                                    NpyAuxData *NPY_UNUSED(auxdata))
+ascii_to_ascii(PyArrayMethod_Context *context, char *const data[],
+               npy_intp const dimensions[], npy_intp const strides[],
+               NpyAuxData *NPY_UNUSED(auxdata))
 {
     PyArray_Descr **descrs = context->descriptors;
     long in_size = ((ASCIIDTypeObject *)descrs[0])->size;
@@ -87,7 +61,7 @@ ascii_to_ascii_strided_or_unaligned(PyArrayMethod_Context *context,
     npy_intp out_stride = strides[1];
 
     while (N--) {
-        memcpy(out, in, out_size * sizeof(char));  // NOLINT
+        memcpy(out, in, copy_size * sizeof(char));  // NOLINT
         for (int i = copy_size; i < out_size; i++) {
             *(out + i) = '\0';
         }
@@ -106,20 +80,7 @@ ascii_to_ascii_get_loop(PyArrayMethod_Context *context, int aligned,
                         NpyAuxData **NPY_UNUSED(out_transferdata),
                         NPY_ARRAYMETHOD_FLAGS *flags)
 {
-    PyArray_Descr **descrs = context->descriptors;
-
-    int contig = (strides[0] == ((ASCIIDTypeObject *)descrs[0])->size *
-                                        sizeof(char) &&
-                  strides[1] == ((ASCIIDTypeObject *)descrs[1])->size *
-                                        sizeof(char));
-
-    if (aligned && contig) {
-        *out_loop = (PyArrayMethod_StridedLoop *)&ascii_to_ascii_contiguous;
-    }
-    else {
-        *out_loop = (PyArrayMethod_StridedLoop
-                             *)&ascii_to_ascii_strided_or_unaligned;
-    }
+    *out_loop = (PyArrayMethod_StridedLoop *)&ascii_to_ascii;
 
     *flags = 0;
     return 0;

asciidtype/asciidtype/src/dtype.c

@@ -113,7 +113,7 @@ asciidtype_setitem(ASCIIDTypeObject *descr, PyObject *obj, char *dataptr)
 
     Py_ssize_t len = PyBytes_Size(value);
 
-    size_t copysize;
+    long copysize;
 
     if (len > descr->size) {
         copysize = descr->size;
@@ -138,7 +138,13 @@ asciidtype_setitem(ASCIIDTypeObject *descr, PyObject *obj, char *dataptr)
 static PyObject *
 asciidtype_getitem(ASCIIDTypeObject *descr, char *dataptr)
 {
-    PyObject *val_obj = PyUnicode_FromString(dataptr);
+    char scalar_buffer[descr->size + 1];
+
+    memcpy(scalar_buffer, dataptr, descr->size * sizeof(char));
+
+    scalar_buffer[descr->size] = '\0';
+
+    PyObject *val_obj = PyUnicode_FromString(scalar_buffer);
     if (val_obj == NULL) {
         return NULL;
     }
@@ -205,7 +211,7 @@ asciidtype_repr(ASCIIDTypeObject *self)
 }
 
 static PyMemberDef ASCIIDType_members[] = {
-        {"size", T_OBJECT_EX, offsetof(ASCIIDTypeObject, size), READONLY,
+        {"size", T_LONG, offsetof(ASCIIDTypeObject, size), READONLY,
          "The number of characters per array element"},
         {NULL},
 };

asciidtype/pyproject.toml

@@ -8,6 +8,9 @@ requires = [
 ]
 build-backend = "mesonpy"
 
+[tool.black]
+line-length = 79
+
 [project]
 name = "asciidtype"
 description = "A dtype for ASCII data"

asciidtype/tests/test_asciidtype.py

@@ -10,11 +10,64 @@ def test_dtype_creation():
 
 def test_scalar_creation():
     dtype = ASCIIDType(7)
-    ASCIIScalar('string', dtype)
+    ASCIIScalar("string", dtype)
 
 
 def test_creation_with_explicit_dtype():
     dtype = ASCIIDType(7)
     arr = np.array(["hello", "this", "is", "an", "array"], dtype=dtype)
     assert repr(arr) == (
-        "array(['hello', 'this', 'is', 'an', 'array'], dtype=ASCIIDType(7))")
+        "array(['hello', 'this', 'is', 'an', 'array'], dtype=ASCIIDType(7))"
+    )
+
+
+def test_creation_truncation():
+    inp = ["hello", "this", "is", "an", "array"]
+
+    dtype = ASCIIDType(5)
+    arr = np.array(inp, dtype=dtype)
+    assert repr(arr) == (
+        "array(['hello', 'this', 'is', 'an', 'array'], dtype=ASCIIDType(5))"
+    )
+
+    dtype = ASCIIDType(4)
+    arr = np.array(inp, dtype=dtype)
+    assert repr(arr) == (
+        "array(['hell', 'this', 'is', 'an', 'arra'], dtype=ASCIIDType(4))"
+    )
+
+    dtype = ASCIIDType(1)
+    arr = np.array(inp, dtype=dtype)
+    assert repr(arr) == (
+        "array(['h', 't', 'i', 'a', 'a'], dtype=ASCIIDType(1))"
+    )
+    assert arr.tobytes() == b"htiaa"
+
+    # dtype = ASCIIDType()
+    # arr = np.array(["hello", "this", "is", "an", "array"], dtype=dtype)
+    # assert repr(arr) == ("array(['', '', '', '', ''], dtype=ASCIIDType(0))")
+    # assert arr.tobytes() == b""
+
+
+def test_casting_to_asciidtype():
+    arr = np.array(["hello", "this", "is", "an", "array"], dtype=ASCIIDType(5))
+
+    assert repr(arr.astype(ASCIIDType(7))) == (
+        "array(['hello', 'this', 'is', 'an', 'array'], dtype=ASCIIDType(7))"
+    )
+
+    assert repr(arr.astype(ASCIIDType(5))) == (
+        "array(['hello', 'this', 'is', 'an', 'array'], dtype=ASCIIDType(5))"
+    )
+
+    assert repr(arr.astype(ASCIIDType(4))) == (
+        "array(['hell', 'this', 'is', 'an', 'arra'], dtype=ASCIIDType(4))"
+    )
+
+    assert repr(arr.astype(ASCIIDType(1))) == (
+        "array(['h', 't', 'i', 'a', 'a'], dtype=ASCIIDType(1))"
+    )
+
+    # assert repr(arr.astype(ASCIIDType())) == (
+    #    "array(['', '', '', '', ''], dtype=ASCIIDType(0))"
+    # )