ENH: f2py: support Fortran size function with two arguments (ticket #1765).

3 files changed, 53 insertions, 1 deletions

diff --git a/numpy/f2py/crackfortran.py b/numpy/f2py/crackfortran.py
index c3082bbf7..aa4c3bc0a 100755
--- a/numpy/f2py/crackfortran.py
+++ b/numpy/f2py/crackfortran.py
@@ -2037,7 +2037,7 @@ def _eval_length(length,params):
         return '(*)'
     return _eval_scalar(length,params)
 
-_is_kind_number = re.compile('\d+_').match
+_is_kind_number = re.compile(r'\d+_').match
 
 def _eval_scalar(value,params):
     if _is_kind_number(value):
@@ -2052,6 +2052,13 @@ 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)
@@ -2194,6 +2201,7 @@ 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]):
@@ -2310,6 +2318,7 @@ 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
new file mode 100644
index 000000000..9602837fe
--- /dev/null
+++ b/numpy/f2py/tests/src/size/foo.f90
@@ -0,0 +1,14 @@
+
+subroutine foo(a, n, m, b)
+  implicit none
+
+  real, intent(in) :: a(n, m)
+  integer, intent(in) :: n, m
+  real, intent(out) :: b(size(a, 1))
+
+  integer :: i
+
+  do i = 1, size(b)
+    b(i) = sum(a(i,:))
+  enddo
+end subroutine
diff --git a/numpy/f2py/tests/test_size.py b/numpy/f2py/tests/test_size.py
new file mode 100644
index 000000000..c00dd9a31
--- /dev/null
+++ b/numpy/f2py/tests/test_size.py
@@ -0,0 +1,29 @@
+import os
+import math
+
+from numpy.testing import *
+from numpy import array
+
+import util
+
+def _path(*a):
+    return os.path.join(*((os.path.dirname(__file__),) + a))
+
+class TestSizeSumExample(util.F2PyTest):
+    sources = [_path('src', 'size', 'foo.f90'),
+               ]
+
+    @dec.slow
+    def test_all(self):
+        r = self.module.foo([[1,2]])
+        assert_equal(r, [3],`r`)
+
+        r = self.module.foo([[1,2],[3,4]])
+        assert_equal(r, [3,7],`r`)
+
+        r = self.module.foo([[1,2],[3,4],[5,6]])
+        assert_equal(r, [3,7,11],`r`)
+
+if __name__ == "__main__":
+    import nose
+    nose.runmodule()