Support functions returning results as an array

Fixes #13

Author: rjfarmer

gfort2py/fArrays.py

@@ -223,6 +223,12 @@ def print(self):
             print(f"\t ubound {self.cvalue.dims[i].ubound}")
             print(f"\t stride {self.cvalue.dims[i].stride}")
 
+    def _make_empty(self):
+        dtype = self.obj.dtype()
+        shape = self.obj.shape()
+
+        return np.zeros(shape, dtype=dtype, order="F")
+
 
 class fAssumedSize(fArray_t):
     def ctype(self):

gfort2py/fProc.py

@@ -8,6 +8,7 @@
 
 from .fVar import fVar
 from .fVar_t import fVar_t
+from .utils import resolve_other_args
 
 _TEST_FLAG = os.environ.get("_GFORT2PY_TEST_FLAG") is not None
 
@@ -72,8 +73,6 @@ def name(self):
         return self.obj.name
 
     def __call__(self, *args, **kwargs):
-        self._set_return()
-
         func_args = self._convert_args(*args, **kwargs)
 
         with _captureStdOut() as cs:
@@ -84,13 +83,13 @@ def __call__(self, *args, **kwargs):
 
         return self._convert_result(res, func_args)
 
-    def _set_return(self):
+    def _set_return(self, other_args):
         if self.obj.is_subroutine():
             self._func.restype = None  # Subroutine
         else:
-            if (
-                self.return_var.obj.is_char()
-            ):  # Returning a character is done as a character + len at start of arg list
+            self.return_var.obj = resolve_other_args(self.return_var.obj, other_args)
+
+            if self.return_var.obj.is_returned_as_arg():
                 self._func.restype = None
             else:
                 self._func.restype = self.return_var.ctype()
@@ -102,6 +101,9 @@ def args_start(self):
                 l = self.return_var.len()
                 res.append(self.return_var.from_param(" " * l))
                 res.append(self.return_var.ctype_len())
+            elif self.return_var.obj.is_always_explicit() and self.obj.is_array():
+                empty_array = self.return_var._make_empty()
+                res.append(ctypes.pointer(self.return_var.from_param(empty_array)))
 
         return res
 
@@ -145,10 +147,13 @@ def args_convert(self, input_args):
         return args, args_end
 
     def _convert_args(self, *args, **kwargs):
-        args_start = self.args_start()
-
         self.input_args = self.args_check(*args, **kwargs)
 
+        # Set this now after arsg_check as we need to be able to
+        # resolve runtime arguments
+        self._set_return(self.input_args)
+        args_start = self.args_start()
+
         args_mid, args_end = self.args_convert(self.input_args)
 
         return args_start + args_mid + args_end
@@ -157,10 +162,14 @@ def _convert_result(self, result, args):
         res = {}
 
         if self.obj.is_function():
-            if self.return_var.obj.is_char():
-                result = args[0]
-                _ = args.pop(0)
-                _ = args.pop(0)  # Twice to pop first and second value
+            if self.return_var.obj.is_returned_as_arg():
+                if self.return_var.obj.is_char():
+                    result = args[0]
+                    _ = args.pop(0)
+                    _ = args.pop(0)  # Twice to pop first and second value
+                elif self.return_var.obj.is_always_explicit() and self.obj.is_array():
+                    result = self.return_var.value
+                    _ = args.pop(0)
 
         if len(self.obj.args()):
             for var in self.input_args:

gfort2py/fVar.py

@@ -37,7 +37,12 @@ def __new__(cls, obj, *args, **kwargs):
                 return fExplicitArr(obj, *args, **kwargs)
             elif obj.is_assumed_size():
                 return fAssumedSize(obj, *args, **kwargs)
-            elif obj.is_assumed_shape() or obj.is_allocatable() or obj.is_pointer():
+            elif (
+                obj.is_assumed_shape()
+                or obj.is_allocatable()
+                or obj.is_pointer()
+                or obj.is_always_explicit()
+            ):
                 return fAssumedShape(obj, *args, **kwargs)
             else:
                 raise TypeError("Unknown array type")

gfort2py/module_parse.py

@@ -254,9 +254,6 @@ def is_function(self):
     def is_array(self):
         return "DIMENSION" in self.sym.attr.attributes
 
-    def is_always_explicit(self):
-        return "ALWAYS_EXPLICIT" in self.sym.attr.attributes
-
     def is_dummy(self):
         return "DUMMY" in self.sym.attr.attributes
 
@@ -269,6 +266,9 @@ def is_return(self):
     def needs_array_desc(self):
         return self.is_dummy() or self.is_allocatable() or self.is_always_explicit()
 
+    def is_always_explicit(self):
+        return "ALWAYS_EXPLICIT" in self.sym.attr.attributes
+
     def not_a_pointer(self):
         return (
             self.needs_array_desc()
@@ -289,6 +289,11 @@ def is_assumed_size(self):
     def is_assumed_shape(self):
         return self.sym.array_spec.array_type == "ASSUMED_SHAPE"
 
+    def is_returned_as_arg(self):
+        return (self.is_always_explicit() and self.is_return() and self.is_array()) or (
+            self.is_char()
+        )
+
     def is_deferred_len(self):
         # Only needed for things that need an extra function argument for their length
         if self.is_char():

tests/explicit_arrays_test.py

@@ -256,12 +256,11 @@ def test_logical_arr_multi(self):
         y = x.func_logical_multi(1.0, 2.0, xarr, 3.0, 4.0)
         self.assertEqual(y.result, True)
 
-    @pytest.mark.skip
     def test_mesh_exp(self):
         # Github issue #13
         i = 5
         y = x.func_mesh_exp(i)
-        np.array_equal(y.result, np.arrange(0, i + 1 + 1))
+        np.array_equal(y.result, np.arange(0, i + 1 + 1))
 
     def test_mesh_exp2(self):
         i = 5
@@ -303,7 +302,16 @@ def test_logical_param_array(self):
             x.const_logical_arr, np.array([True, False, True, False, True], dtype=bool)
         )
 
-    @pytest.mark.skip
     def test_func_return_1d_int_arr(self):
         res = x.func_return_1d_int_arr()
         assert np.array_equal(res.result, np.array([1, 2, 3, 4, 5]))
+
+    def test_func_return_1d_int_arr_n(self):
+        res = x.func_return_1d_int_arr_n(5)
+        assert np.array_equal(res.result, np.array([1, 2, 3, 4, 5]))
+
+    def test_func_return_2d_int_arr(self):
+        res = x.func_return_2d_int_arr()
+        assert np.array_equal(
+            res.result, np.array([1, 2, 3, 4, 5, 6]).reshape(3, 2, order="F")
+        )