1 files changed, 116 insertions, 0 deletions
diff --git a/numpy/f2py/tests/test_crackfortran.py b/numpy/f2py/tests/test_crackfortran.py
index 140f42cbc..039e085b4 100644
--- a/numpy/f2py/tests/test_crackfortran.py
+++ b/numpy/f2py/tests/test_crackfortran.py
@@ -1,3 +1,4 @@
+import pytest
 import numpy as np
 from numpy.testing import assert_array_equal, assert_equal
 from numpy.f2py.crackfortran import markinnerspaces
@@ -39,6 +40,7 @@ class TestNoSpace(util.F2PyTest):
 
 
 class TestPublicPrivate():
+
     def test_defaultPrivate(self, tmp_path):
         f_path = tmp_path / "mod.f90"
         with f_path.open('w') as ff:
@@ -165,3 +167,117 @@ class TestMarkinnerspaces():
     def test_multiple_relevant_spaces(self):
         assert_equal(markinnerspaces("a 'b c' 'd e'"), "a 'b@_@c' 'd@_@e'")
         assert_equal(markinnerspaces(r'a "b c" "d e"'), r'a "b@_@c" "d@_@e"')
+
+
+class TestDimSpec(util.F2PyTest):
+    """This test suite tests various expressions that are used as dimension
+    specifications.
+
+    There exists two usage cases where analyzing dimensions
+    specifications are important.
+
+    In the first case, the size of output arrays must be defined based
+    on the inputs to a Fortran function. Because Fortran supports
+    arbitrary bases for indexing, for instance, `arr(lower:upper)`,
+    f2py has to evaluate an expression `upper - lower + 1` where
+    `lower` and `upper` are arbitrary expressions of input parameters.
+    The evaluation is performed in C, so f2py has to translate Fortran
+    expressions to valid C expressions (an alternative approach is
+    that a developer specifies the corresponding C expressions in a
+    .pyf file).
+
+    In the second case, when user provides an input array with a given
+    size but some hidden parameters used in dimensions specifications
+    need to be determined based on the input array size. This is a
+    harder problem because f2py has to solve the inverse problem: find
+    a parameter `p` such that `upper(p) - lower(p) + 1` equals to the
+    size of input array. In the case when this equation cannot be
+    solved (e.g. because the input array size is wrong), raise an
+    error before calling the Fortran function (that otherwise would
+    likely crash Python process when the size of input arrays is
+    wrong). f2py currently supports this case only when the equation
+    is linear with respect to unknown parameter.
+
+    """
+
+    suffix = '.f90'
+
+    code_template = textwrap.dedent("""
+      function get_arr_size_{count}(a, n) result (length)
+        integer, intent(in) :: n
+        integer, dimension({dimspec}), intent(out) :: a
+        integer length
+        length = size(a)
+      end function
+
+      subroutine get_inv_arr_size_{count}(a, n)
+        integer :: n
+        ! the value of n is computed in f2py wrapper
+        !f2py intent(out) n
+        integer, dimension({dimspec}), intent(in) :: a
+        if (a({first}).gt.0) then
+          print*, "a=", a
+        endif
+      end subroutine
+    """)
+
+    linear_dimspecs = ['n', '2*n', '2:n', 'n/2', '5 - n/2', '3*n:20',
+                       'n*(n+1):n*(n+5)']
+    nonlinear_dimspecs = ['2*n:3*n*n+2*n']
+    all_dimspecs = linear_dimspecs + nonlinear_dimspecs
+
+    code = ''
+    for count, dimspec in enumerate(all_dimspecs):
+        code += code_template.format(
+            count=count, dimspec=dimspec,
+            first=dimspec.split(':')[0] if ':' in dimspec else '1')
+
+    @pytest.mark.parametrize('dimspec', all_dimspecs)
+    def test_array_size(self, dimspec):
+
+        count = self.all_dimspecs.index(dimspec)
+        get_arr_size = getattr(self.module, f'get_arr_size_{count}')
+
+        for n in [1, 2, 3, 4, 5]:
+            sz, a = get_arr_size(n)
+            assert len(a) == sz
+
+    @pytest.mark.parametrize('dimspec', all_dimspecs)
+    def test_inv_array_size(self, dimspec):
+
+        count = self.all_dimspecs.index(dimspec)
+        get_arr_size = getattr(self.module, f'get_arr_size_{count}')
+        get_inv_arr_size = getattr(self.module, f'get_inv_arr_size_{count}')
+
+        for n in [1, 2, 3, 4, 5]:
+            sz, a = get_arr_size(n)
+            if dimspec in self.nonlinear_dimspecs:
+                # one must specify n as input, the call we'll ensure
+                # that a and n are compatible:
+                n1 = get_inv_arr_size(a, n)
+            else:
+                # in case of linear dependence, n can be determined
+                # from the shape of a:
+                n1 = get_inv_arr_size(a)
+            # n1 may be different from n (for instance, when `a` size
+            # is a function of some `n` fraction) but it must produce
+            # the same sized array
+            sz1, _ = get_arr_size(n1)
+            assert sz == sz1, (n, n1, sz, sz1)
+
+
+class TestModuleDeclaration():
+    def test_dependencies(self, tmp_path):
+        f_path = tmp_path / "mod.f90"
+        with f_path.open('w') as ff:
+            ff.write(textwrap.dedent("""\
+            module foo
+              type bar
+                character(len = 4) :: text
+              end type bar
+              type(bar), parameter :: abar = bar('abar')
+            end module foo
+            """))
+        mod = crackfortran.crackfortran([str(f_path)])
+        assert len(mod) == 1
+        assert mod[0]['vars']['abar']['='] == "bar('abar')"