1 files changed, 30 insertions, 1 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))'