Merge pull request #18759 from pearu/gh-18431-string_from_pyobj

BUG: revise string_from_pyobj/try_pyarr_from_string with respect to malloc and copy.

5 files changed, 252 insertions, 59 deletions

diff --git a/numpy/f2py/cfuncs.py b/numpy/f2py/cfuncs.py
index f403a66b5..60685be8a 100644
--- a/numpy/f2py/cfuncs.py
+++ b/numpy/f2py/cfuncs.py
@@ -469,7 +469,7 @@ cppmacros['MEMCOPY'] = """\
 """
 cppmacros['STRINGMALLOC'] = """\
 #define STRINGMALLOC(str,len)\\
-    if ((str = (string)malloc(sizeof(char)*(len+1))) == NULL) {\\
+    if ((str = (string)malloc(len+1)) == NULL) {\\
         PyErr_SetString(PyExc_MemoryError, \"out of memory\");\\
         goto capi_fail;\\
     } else {\\
@@ -481,18 +481,18 @@ cppmacros['STRINGFREE'] = """\
 """
 needs['STRINGCOPYN'] = ['string.h', 'FAILNULL']
 cppmacros['STRINGCOPYN'] = """\
-#define STRINGCOPYN(to,from,buf_size)                           \\
+/*
+STRINGCOPYN copies N bytes.
+
+`to` and `from` buffers must have sizes of at least N bytes.
+*/
+#define STRINGCOPYN(to,from,N)                                  \\
     do {                                                        \\
-        int _m = (buf_size);                                    \\
+        int _m = (N);                                           \\
         char *_to = (to);                                       \\
         char *_from = (from);                                   \\
         FAILNULL(_to); FAILNULL(_from);                         \\
-        (void)strncpy(_to, _from, sizeof(char)*_m);             \\
-        _to[_m-1] = '\\0';                                      \\
-        /* Padding with spaces instead of nulls */              \\
-        for (_m -= 2; _m >= 0 && _to[_m] == '\\0'; _m--) {      \\
-            _to[_m] = ' ';                                      \\
-        }                                                       \\
+        (void)strncpy(_to, _from, _m);             \\
     } while (0)
 """
 needs['STRINGCOPY'] = ['string.h', 'FAILNULL']
@@ -623,71 +623,121 @@ static int *nextforcomb(void) {
 }"""
 needs['try_pyarr_from_string'] = ['STRINGCOPYN', 'PRINTPYOBJERR', 'string']
 cfuncs['try_pyarr_from_string'] = """\
-static int try_pyarr_from_string(PyObject *obj,const string str) {
-    PyArrayObject *arr = NULL;
-    if (PyArray_Check(obj) && (!((arr = (PyArrayObject *)obj) == NULL)))
-        { STRINGCOPYN(PyArray_DATA(arr),str,PyArray_NBYTES(arr)); }
-    return 1;
+/*
+  try_pyarr_from_string copies str[:len(obj)] to the data of an `ndarray`.
+
+  If obj is an `ndarray`, it is assumed to be contiguous.
+
+  If the specified len==-1, str must be null-terminated.
+*/
+static int try_pyarr_from_string(PyObject *obj,
+                                 const string str, const int len) {
+#ifdef DEBUGCFUNCS
+fprintf(stderr, "try_pyarr_from_string(str='%s', len=%d, obj=%p)\\n",
+        (char*)str,len, obj);
+#endif
+    if (PyArray_Check(obj)) {
+        PyArrayObject *arr = (PyArrayObject *)obj;
+        assert(ISCONTIGUOUS(arr));
+        string buf = PyArray_DATA(arr);
+        npy_intp n = len;
+        if (n == -1) {
+            /* Assuming null-terminated str. */
+            n = strlen(str);
+        }
+        if (n > PyArray_NBYTES(arr)) {
+            n = PyArray_NBYTES(arr);
+        }
+        STRINGCOPYN(buf, str, n);
+        return 1;
+    }
 capi_fail:
     PRINTPYOBJERR(obj);
-    PyErr_SetString(#modulename#_error,\"try_pyarr_from_string failed\");
+    PyErr_SetString(#modulename#_error, \"try_pyarr_from_string failed\");
     return 0;
 }
 """
 needs['string_from_pyobj'] = ['string', 'STRINGMALLOC', 'STRINGCOPYN']
 cfuncs['string_from_pyobj'] = """\
+/*
+  Create a new string buffer `str` of at most length `len` from a
+  Python string-like object `obj`.
+
+  The string buffer has given size (len) or the size of inistr when len==-1.
+
+  The string buffer is null-terminated.
+ */
 static int
-string_from_pyobj(string *str,int *len,const string inistr,PyObject *obj,const char *errmess)
+string_from_pyobj(string *str, int *len, const string inistr, PyObject *obj,
+                  const char *errmess)
 {
-    PyArrayObject *arr = NULL;
     PyObject *tmp = NULL;
+    string buf = NULL;
+    npy_intp n = -1;
 #ifdef DEBUGCFUNCS
-fprintf(stderr,\"string_from_pyobj(str='%s',len=%d,inistr='%s',obj=%p)\\n\",(char*)str,*len,(char *)inistr,obj);
+fprintf(stderr,\"string_from_pyobj(str='%s',len=%d,inistr='%s',obj=%p)\\n\",
+               (char*)str, *len, (char *)inistr, obj);
 #endif
     if (obj == Py_None) {
-        if (*len == -1)
-            *len = strlen(inistr); /* Will this cause problems? */
-        STRINGMALLOC(*str,*len);
-        STRINGCOPYN(*str,inistr,*len+1);
-        return 1;
+        n = strlen(inistr);
+        buf = inistr;
     }
-    if (PyArray_Check(obj)) {
-        if ((arr = (PyArrayObject *)obj) == NULL)
-            goto capi_fail;
+    else if (PyArray_Check(obj)) {
+        PyArrayObject *arr = (PyArrayObject *)obj;
         if (!ISCONTIGUOUS(arr)) {
-            PyErr_SetString(PyExc_ValueError,\"array object is non-contiguous.\");
+            PyErr_SetString(PyExc_ValueError,
+                            \"array object is non-contiguous.\");
             goto capi_fail;
         }
-        if (*len == -1)
-            *len = (PyArray_ITEMSIZE(arr))*PyArray_SIZE(arr);
-        STRINGMALLOC(*str,*len);
-        STRINGCOPYN(*str,PyArray_DATA(arr),*len+1);
-        return 1;
-    }
-    if (PyBytes_Check(obj)) {
-        tmp = obj;
-        Py_INCREF(tmp);
-    }
-    else if (PyUnicode_Check(obj)) {
-        tmp = PyUnicode_AsASCIIString(obj);
+        n = PyArray_NBYTES(arr);
+        buf = PyArray_DATA(arr);
     }
     else {
-        PyObject *tmp2;
-        tmp2 = PyObject_Str(obj);
-        if (tmp2) {
-            tmp = PyUnicode_AsASCIIString(tmp2);
-            Py_DECREF(tmp2);
+        if (PyBytes_Check(obj)) {
+            tmp = obj;
+            Py_INCREF(tmp);
+        }
+        else if (PyUnicode_Check(obj)) {
+            tmp = PyUnicode_AsASCIIString(obj);
         }
         else {
-            tmp = NULL;
+            PyObject *tmp2;
+            tmp2 = PyObject_Str(obj);
+            if (tmp2) {
+                tmp = PyUnicode_AsASCIIString(tmp2);
+                Py_DECREF(tmp2);
+            }
+            else {
+                tmp = NULL;
+            }
         }
+        if (tmp == NULL) goto capi_fail;
+        n = PyBytes_GET_SIZE(tmp);
+        buf = PyBytes_AS_STRING(tmp);
     }
-    if (tmp == NULL) goto capi_fail;
-    if (*len == -1)
-        *len = PyBytes_GET_SIZE(tmp);
-    STRINGMALLOC(*str,*len);
-    STRINGCOPYN(*str,PyBytes_AS_STRING(tmp),*len+1);
-    Py_DECREF(tmp);
+    if (*len == -1) {
+        /* TODO: change the type of `len` so that we can remove this */
+        if (n > NPY_MAX_INT) {
+            PyErr_SetString(PyExc_OverflowError,
+                            "object too large for a 32-bit int");
+            goto capi_fail;
+        }
+        *len = n;
+    }
+    else if (*len < n) {
+        /* discard the last (len-n) bytes of input buf */
+        n = *len;  
+    }
+    if (n < 0 || *len < 0 || buf == NULL) {
+        goto capi_fail;
+    }
+    STRINGMALLOC(*str, *len);  // *str is allocated with size (*len + 1)
+    if (n < *len) {
+        /* Pad fixed-width string with nulls */
+        memset(*str + n, '\\0', *len - n);
+    }
+    STRINGCOPYN(*str, buf, n);
+    Py_XDECREF(tmp);
     return 1;
 capi_fail:
     Py_XDECREF(tmp);
@@ -702,7 +752,6 @@ capi_fail:
 }
 """
 
-
 needs['char_from_pyobj'] = ['int_from_pyobj']
 cfuncs['char_from_pyobj'] = """\
 static int
diff --git a/numpy/f2py/rules.py b/numpy/f2py/rules.py
index 63e47baa2..f01b2bcfa 100755
--- a/numpy/f2py/rules.py
+++ b/numpy/f2py/rules.py
@@ -561,7 +561,8 @@ rout_rules = [
                  '\tint #name#_return_value_len = 0;'],
         'callfortran':'#name#_return_value,#name#_return_value_len,',
         'callfortranroutine':['\t#name#_return_value_len = #rlength#;',
-                              '\tif ((#name#_return_value = (string)malloc(sizeof(char)*(#name#_return_value_len+1))) == NULL) {',
+                              '\tif ((#name#_return_value = (string)malloc('
+                                  '#name#_return_value_len+1) == NULL) {',
                               '\t\tPyErr_SetString(PyExc_MemoryError, \"out of memory\");',
                               '\t\tf2py_success = 0;',
                               '\t} else {',
@@ -963,7 +964,8 @@ if (#varname#_cb.capi==Py_None) {
         'args_capi': {isrequired: ',&#varname#_capi'},
         'keys_capi': {isoptional: ',&#varname#_capi'},
         'pyobjfrom': {isintent_inout: '''\
-\tf2py_success = try_pyarr_from_#ctype#(#varname#_capi,#varname#);
+\tf2py_success = try_pyarr_from_#ctype#(#varname#_capi, #varname#,
+\t                                      slen(#varname#));
 \tif (f2py_success) {'''},
         'closepyobjfrom': {isintent_inout: '\t} /*if (f2py_success) of #varname# pyobjfrom*/'},
         'need': {isintent_inout: 'try_pyarr_from_#ctype#'},
diff --git a/numpy/f2py/tests/src/array_from_pyobj/wrapmodule.c b/numpy/f2py/tests/src/array_from_pyobj/wrapmodule.c
index 0411b62e0..fe21d4b9b 100644
--- a/numpy/f2py/tests/src/array_from_pyobj/wrapmodule.c
+++ b/numpy/f2py/tests/src/array_from_pyobj/wrapmodule.c
@@ -93,7 +93,7 @@ static PyObject *f2py_rout_wrap_attrs(PyObject *capi_self,
   PyObject *strides = NULL;
   char s[100];
   int i;
-  memset(s,0,100*sizeof(char));
+  memset(s,0,100);
   if (!PyArg_ParseTuple(capi_args,"O!|:wrap.attrs",
                         &PyArray_Type,&arr_capi))
     return NULL;
diff --git a/numpy/f2py/tests/test_return_character.py b/numpy/f2py/tests/test_return_character.py
index 429e69bb4..dc524e60c 100644
--- a/numpy/f2py/tests/test_return_character.py
+++ b/numpy/f2py/tests/test_return_character.py
@@ -24,8 +24,8 @@ class TestReturnCharacter(util.F2PyTest):
             assert_(t(23) == b'23        ', repr(t(23)))
             assert_(t('123456789abcdef') == b'123456789a')
         elif tname in ['t5', 's5']:
-            assert_(t(23) == b'23   ', repr(t(23)))
-            assert_(t('ab') == b'ab   ', repr(t('ab')))
+            assert_(t(23) == b'23', repr(t(23)))
+            assert_(t('ab') == b'ab', repr(t('ab')))
             assert_(t('123456789abcdef') == b'12345')
         else:
             raise NotImplementedError
diff --git a/numpy/f2py/tests/test_string.py b/numpy/f2py/tests/test_string.py
index e3ec96af9..1530c58f6 100644
--- a/numpy/f2py/tests/test_string.py
+++ b/numpy/f2py/tests/test_string.py
@@ -1,6 +1,6 @@
 import os
 import pytest
-
+import textwrap
 from numpy.testing import assert_array_equal
 import numpy as np
 from . import util
@@ -9,14 +9,156 @@ from . import util
 def _path(*a):
     return os.path.join(*((os.path.dirname(__file__),) + a))
 
+
 class TestString(util.F2PyTest):
     sources = [_path('src', 'string', 'char.f90')]
 
     @pytest.mark.slow
     def test_char(self):
         strings = np.array(['ab', 'cd', 'ef'], dtype='c').T
-        inp, out = self.module.char_test.change_strings(strings, strings.shape[1])
+        inp, out = self.module.char_test.change_strings(strings,
+                                                        strings.shape[1])
         assert_array_equal(inp, strings)
         expected = strings.copy()
         expected[1, :] = 'AAA'
         assert_array_equal(out, expected)
+
+
+class TestDocStringArguments(util.F2PyTest):
+    suffix = '.f'
+
+    code = """
+C FILE: STRING.F
+      SUBROUTINE FOO(A,B,C,D)
+      CHARACTER*5 A, B
+      CHARACTER*(*) C,D
+Cf2py intent(in) a,c
+Cf2py intent(inout) b,d
+      PRINT*, "A=",A
+      PRINT*, "B=",B
+      PRINT*, "C=",C
+      PRINT*, "D=",D
+      PRINT*, "CHANGE A,B,C,D"
+      A(1:1) = 'A'
+      B(1:1) = 'B'
+      C(1:1) = 'C'
+      D(1:1) = 'D'
+      PRINT*, "A=",A
+      PRINT*, "B=",B
+      PRINT*, "C=",C
+      PRINT*, "D=",D
+      END
+C END OF FILE STRING.F
+        """
+
+    def test_example(self):
+        a = np.array(b'123\0\0')
+        b = np.array(b'123\0\0')
+        c = np.array(b'123')
+        d = np.array(b'123')
+
+        self.module.foo(a, b, c, d)
+
+        assert a.tobytes() == b'123\0\0'
+        assert b.tobytes() == b'B23\0\0', (b.tobytes(),)
+        assert c.tobytes() == b'123'
+        assert d.tobytes() == b'D23'
+
+
+class TestFixedString(util.F2PyTest):
+    suffix = '.f90'
+
+    code = textwrap.dedent("""
+       function sint(s) result(i)
+          implicit none
+          character(len=*) :: s
+          integer :: j, i
+          i = 0
+          do j=1, len(s)
+           i = i + ichar(s(j:j)) * 10 ** (j - 1)
+          end do
+          return
+        end function sint
+
+        function test_in_bytes4(a) result (i)
+          implicit none
+          integer :: sint
+          character(len=4) :: a
+          integer :: i
+          i = sint(a)
+          a(1:1) = 'A'
+          return
+        end function test_in_bytes4
+
+        function test_inout_bytes4(a) result (i)
+          implicit none
+          integer :: sint
+          character(len=4), intent(inout) :: a
+          integer :: i
+          if (ichar(a(1:1)).ne.0) then
+            a(1:1) = 'E'
+          endif
+          i = sint(a)
+          return
+        end function test_inout_bytes4
+        """)
+
+    @staticmethod
+    def _sint(s, start=0, end=None):
+        """Return the content of a string buffer as integer value.
+
+        For example:
+          _sint('1234') -> 4321
+          _sint('123A') -> 17321
+        """
+        if isinstance(s, np.ndarray):
+            s = s.tobytes()
+        elif isinstance(s, str):
+            s = s.encode()
+        assert isinstance(s, bytes)
+        if end is None:
+            end = len(s)
+        i = 0
+        for j in range(start, min(end, len(s))):
+            i += s[j] * 10 ** j
+        return i
+
+    def _get_input(self, intent='in'):
+        if intent in ['in']:
+            yield ''
+            yield '1'
+            yield '1234'
+            yield '12345'
+            yield b''
+            yield b'\0'
+            yield b'1'
+            yield b'\01'
+            yield b'1\0'
+            yield b'1234'
+            yield b'12345'
+        yield np.ndarray((), np.bytes_, buffer=b'')  # array(b'', dtype='|S0')
+        yield np.array(b'')                          # array(b'', dtype='|S1')
+        yield np.array(b'\0')
+        yield np.array(b'1')
+        yield np.array(b'1\0')
+        yield np.array(b'\01')
+        yield np.array(b'1234')
+        yield np.array(b'123\0')
+        yield np.array(b'12345')
+
+    def test_intent_in(self):
+        for s in self._get_input():
+            r = self.module.test_in_bytes4(s)
+            # also checks that s is not changed inplace
+            expected = self._sint(s, end=4)
+            assert r == expected, (s)
+
+    def test_intent_inout(self):
+        for s in self._get_input(intent='inout'):
+            rest = self._sint(s, start=4)
+            r = self.module.test_inout_bytes4(s)
+            expected = self._sint(s, end=4)
+            assert r == expected
+
+            # check that the rest of input string is preserved
+            assert rest == self._sint(s, start=4)