BUG: Fix two argument size support for Fortran module routines. Reverted size-to-shape mapping patch and implemented two argument size function in C.

4 files changed, 76 insertions, 10 deletions

diff --git a/numpy/f2py/cfuncs.py b/numpy/f2py/cfuncs.py
index 56a193963..99515b42b 100644
--- a/numpy/f2py/cfuncs.py
+++ b/numpy/f2py/cfuncs.py
@@ -59,6 +59,7 @@ includes['arrayobject.h']='''#define PY_ARRAY_UNIQUE_SYMBOL PyArray_API
 #include "arrayobject.h"'''
 
 includes['arrayobject.h']='#include "fortranobject.h"'
+includes['stdarg.h']='#include <stdarg.h>'
 
 ############# Type definitions ###############
 
@@ -243,6 +244,7 @@ cppmacros['MINMAX']="""\
 #define MIN(a,b) ((a < b) ? (a) : (b))
 #endif
 """
+needs['len..']=['f2py_size']
 cppmacros['len..']="""\
 #define rank(var) var ## _Rank
 #define shape(var,dim) var ## _Dims[dim]
@@ -251,9 +253,36 @@ cppmacros['len..']="""\
 #define fshape(var,dim) shape(var,rank(var)-dim-1)
 #define len(var) shape(var,0)
 #define flen(var) fshape(var,0)
-#define size(var) PyArray_SIZE((PyArrayObject *)(capi_ ## var ## _tmp))
+#define old_size(var) PyArray_SIZE((PyArrayObject *)(capi_ ## var ## _tmp))
 /* #define index(i) capi_i ## i */
 #define slen(var) capi_ ## var ## _len
+#define size(var, dim...) f2py_size((PyArrayObject *)(capi_ ## var ## _tmp), ##dim, -1)
+"""
+needs['f2py_size']=['stdarg.h']
+cfuncs['f2py_size']="""\
+int f2py_size(PyArrayObject* var, ...)
+{
+  npy_int sz = 0;
+  npy_int dim;
+  npy_int rank;
+  va_list argp;
+  va_start(argp, var);
+  dim = va_arg(argp, npy_int);
+  if (dim==-1)
+    {
+      sz = PyArray_SIZE(var);
+    }
+  else
+    {
+      rank = PyArray_NDIM(var);
+      if (dim>=1 && dim<=rank)
+        sz = PyArray_DIM(var, dim-1);
+      else
+        fprintf(stderr, \"f2py_size: 2nd argument value=%d fails to satisfy 1<=value<=%d. Result will be 0.\\n\", dim, rank);
+    }
+  va_end(argp);
+  return sz;
+}
 """
 
 cppmacros['pyobj_from_char1']='#define pyobj_from_char1(v) (PyInt_FromLong(v))'
diff --git a/numpy/f2py/crackfortran.py b/numpy/f2py/crackfortran.py
index 8299e0aa6..6292bdd1a 100755
--- a/numpy/f2py/crackfortran.py
+++ b/numpy/f2py/crackfortran.py
@@ -2058,13 +2058,6 @@ def _eval_scalar(value,params):
                 % (msg,value,params.keys()))
     return value
 
-_size_call_sub = re.compile(r'size\s*\((?P<arg1>\w+)\s*[,]').sub
-def two_argument_size_hook(expr):
-    new_expr = _size_call_sub(r'shape(\g<arg1>,-1+', expr)
-    if verbose > 1 and expr!=new_expr:
-        outmess('two_argument_size_hook: mapping %r to %r\n' % (expr, new_expr))
-    return new_expr
-
 def analyzevars(block):
     global f90modulevars
     setmesstext(block)
@@ -2207,7 +2200,6 @@ def analyzevars(block):
                         if d[:4] == '1 * ': d = d[4:]
                         if di and di[-4:] == '/(1)': di = di[:-4]
                         if v: savelindims[d] = v,di
-                    d = two_argument_size_hook(d)
                     vars[n]['dimension'].append(d)
         if 'dimension' in vars[n]:
             if isintent_c(vars[n]):
@@ -2324,7 +2316,6 @@ def analyzevars(block):
                 if not vars[n]['depend']: del vars[n]['depend']
             if isscalar(vars[n]):
                 vars[n]['='] = _eval_scalar(vars[n]['='],params)
-                vars[n]['='] = two_argument_size_hook(vars[n]['='])
 
     for n in vars.keys():
         if n==block['name']: # n is block name
diff --git a/numpy/f2py/tests/src/size/foo.f90 b/numpy/f2py/tests/src/size/foo.f90
index 9602837fe..5b66f8c43 100644
--- a/numpy/f2py/tests/src/size/foo.f90
+++ b/numpy/f2py/tests/src/size/foo.f90
@@ -12,3 +12,33 @@ subroutine foo(a, n, m, b)
     b(i) = sum(a(i,:))
   enddo
 end subroutine
+
+subroutine trans(x,y)
+  implicit none
+  real, intent(in), dimension(:,:) :: x
+  real, intent(out), dimension( size(x,2), size(x,1) ) :: y
+  integer :: N, M, i, j
+  N = size(x,1)
+  M = size(x,2)
+  DO i=1,N
+     do j=1,M
+        y(j,i) = x(i,j)
+     END DO
+  END DO
+end subroutine trans
+
+subroutine flatten(x,y)
+  implicit none
+  real, intent(in), dimension(:,:) :: x
+  real, intent(out), dimension( size(x) ) :: y
+  integer :: N, M, i, j, k
+  N = size(x,1)
+  M = size(x,2)
+  k = 1
+  DO i=1,N
+     do j=1,M
+        y(k) = x(i,j)
+        k = k + 1
+     END DO
+  END DO
+end subroutine flatten
diff --git a/numpy/f2py/tests/test_size.py b/numpy/f2py/tests/test_size.py
index c00dd9a31..a548e9885 100644
--- a/numpy/f2py/tests/test_size.py
+++ b/numpy/f2py/tests/test_size.py
@@ -24,6 +24,22 @@ class TestSizeSumExample(util.F2PyTest):
         r = self.module.foo([[1,2],[3,4],[5,6]])
         assert_equal(r, [3,7,11],`r`)
 
+    @dec.slow
+    def test_transpose(self):
+        r = self.module.trans([[1,2]])
+        assert_equal(r, [[1],[2]],`r`)
+
+        r = self.module.trans([[1,2,3],[4,5,6]])
+        assert_equal(r, [[1,4],[2,5],[3,6]],`r`)
+
+    @dec.slow
+    def test_flatten(self):
+        r = self.module.flatten([[1,2]])
+        assert_equal(r, [1,2],`r`)
+
+        r = self.module.flatten([[1,2,3],[4,5,6]])
+        assert_equal(r, [1,2,3,4,5,6],`r`)
+
 if __name__ == "__main__":
     import nose
     nose.runmodule()