feat[cartesian]: lazy annotated gtscript functions

docs/development/ADRs/cartesian/frontend/lazy-functions.md

@@ -0,0 +1,66 @@
+# Lazy annotated gtscript functions
+
+In the context of extending `gt4py.cartesian` with run-time dependent type annotations, we needed a way to delay the annotation of `gtscript.function`s and inject code before that happens. We decided to extend the `gtscript` language with a `lazy_function(before_annotation=None, after_annotation=None)` decorator allowing us to do so.
+
+## Context
+
+NDSL is extending `gt4py.cartesian` with run-time dependent type annotations. We would like NDSL-users to be able to use those type hints in `gtscript.stencil` and `gtscript.function` definitions. To do so, we have a registry in place that lets NDSL-users declare/register run-time dependent types. The declaration results in intermediate "mockup types", which can be resolved to the real type after it has been registered.
+
+The above declare/register approach works well for stencils because we deny users the possibility to use the `@stencil` decorator. Instead, we force usage of a "stencil factory" that allows us to inject the real types before sending stencils to gt4py for annotation.
+
+For `gtscript.function`s, we can't inject the real types because `@gtscript.function` decorators are executed a python parsing time and at this time we don't have the real type information yet. This results in the inconsistent behavior that NDSL-users can use run-time dependent type annotations in stencils but not in functions.
+
+## Decision
+
+We decided to extend the `gtscript` language with a `@gtscript.lazy_function(before_annotation=None, after_annotation=None)` allowing user extending `gt4py.cartesian` to delay the annotation process of those functions until they are actually used from within a stencil or another function.
+
+## Consequences
+
+This change expands the `gtscript` language, which comes with maintenance overhead and allows power users of `gt4py.cartesian` to intercept function annotations before and after annotation happens.
+
+In NDSL, this allows us to default to `gtscript.lazy_functions(before_annotation=resolve_deferred_types)` allowing NDSL to inject code to automatically resolve deferred types. A basic example follows:
+
+```py
+class DeferredType:
+    """This is a mockup type - to be resolved to a "real type" later."""
+    pass
+
+
+def resolve_type(func: Callable):
+    """This function resolves the deferred type"""
+    for name, type_ in func.__annotations__.items():
+        if type_ == DeferredType:
+            func.__annotations__[name] = gtscript.Field[float]
+
+
+# configure lazy_function to resolve the deferred type on the `plus_one()` function
+@gtscript.lazy_function(before_annotation=resolve_type)
+def plus_one(a: DeferredType):
+    return a + 1.0
+
+
+def definition_func(in_field: DeferredType, out_field: gtscript.Field[float]):
+    with computation(PARALLEL), interval(...):
+        out_field = plus_one(in_field)
+
+
+# Manually resolve the deferred type on stencil,
+resolve_type(definition_func)
+
+# then parse the stencil, which will automatically resolve the deferred type in `plus_one()`
+stencil = parse_definition(definition_func)
+```
+
+## Alternatives considered
+
+### Avoid using run-time dependent types
+
+Avoiding types that are run-time dependent solves many issues. It is, however, the given scope of this problem. One could forbid usage of run-time dependent type annotations in `gtscript.function`s. In NDSL, this is, however, inconsistent with the seemingly allowed usage in stencil code and in general forces more mental load to `gt4py` users by having to mentally map untyped function parameters.
+
+### Extend the current `gtscript.function`
+
+It might be possible to extend the existing `gtscript.function` decorator with a flag to distinguish eager and delayed execution. One could argue that this would require less duplicated code. However, I'd argue the bulk of the work is adding support in the frontend, which needs to happen in either case.
+
+## References
+
+Related NDSL issue: <https://github.com/NOAA-GFDL/NDSL/issues/504>.

src/gt4py/cartesian/frontend/gtscript_frontend.py

@@ -2312,6 +2312,16 @@ def collect_external_symbols(definition):
                         loc=nodes.Location.from_ast_node(name_nodes[collected_name][0]),
                     )
 
+        for key, value in nonlocal_symbols.items():
+            # Support @lazy_function() decorators by only evaluating them at this point
+            if (
+                callable(value)
+                and not hasattr(value, "_gtscript_")
+                and value.__qualname__.startswith("lazy_function.")
+            ):
+                value = value()
+                nonlocal_symbols[key] = value
+
         return nonlocal_symbols, imported_symbols
 
     @staticmethod

src/gt4py/cartesian/gtscript.py

@@ -159,14 +159,53 @@ def _parse_annotation(arg, annotation):
     return original_annotations
 
 
-def function(func):
+def function(func: Callable) -> Callable:
     """Mark a GTScript function."""
     from gt4py.cartesian.frontend import gtscript_frontend as gt_frontend
 
     gt_frontend.GTScriptParser.annotate_definition(func)
     return func
 
 
+def lazy_function(
+    *, before_annotation: Callable | None = None, after_annotation: Callable | None = None
+) -> Callable:
+    """
+    Mark a GTScript function that will only be annotated before usage.
+
+    Lazy functions can be useful if e.g. annotated types aren't known when python
+    parses the code. Hooks allow to insert code before and after the function is
+    annotated.
+
+    Parameters
+    ----------
+        before_annotation : Callable | None
+            Hook to be called before annotating the function. Takes the function
+            that is about to be annotated as an argument and returns nothing.
+        after_annotation : Callable | None
+            Hook to be called after annotating the function. Takes the annotated
+            function as an argument and returns nothing.
+    """
+
+    def wrapper(func: Callable) -> Callable:
+        def inner_function() -> Callable:
+            if before_annotation is not None:
+                before_annotation(func)
+
+            from gt4py.cartesian.frontend import gtscript_frontend as gt_frontend
+
+            gt_frontend.GTScriptParser.annotate_definition(func)
+
+            if after_annotation is not None:
+                after_annotation(func)
+
+            return func
+
+        return inner_function
+
+    return wrapper
+
+
 # Interface functions
 def stencil(
     backend,

tests/cartesian_tests/unit_tests/frontend_tests/test_gtscript_frontend.py

@@ -481,6 +481,58 @@ def definition_func(phi: gtscript.Field[np.float64]):
             )
 
 
+class TestLazyFunction:
+    def test_simple_case(self) -> None:
+        @gtscript.lazy_function()
+        def constant():
+            return 1.0
+
+        def definition_func(out_field: gtscript.Field[float]):
+            with computation(PARALLEL), interval(...):
+                out_field = constant()
+
+        stencil = parse_definition(
+            definition_func,
+            name=inspect.stack()[0][3],
+            module=self.__class__.__name__,
+        )
+
+        # check that constant was properly inlined
+        assert stencil.computations[0].body.stmts[0].value.value == 1.0
+
+    def test_annotation_deferred_type(self) -> None:
+        class DeferredType:
+            pass
+
+        def resolve_type(func: Callable):
+            """This function resolves the deferred type"""
+            for name, type_ in func.__annotations__.items():
+                if type_ == DeferredType:
+                    func.__annotations__[name] = gtscript.Field[float]
+
+        # configure lazy_function to resolve the deferred type on the `plus_one()` function
+        @gtscript.lazy_function(before_annotation=resolve_type)
+        def plus_one(a: DeferredType):
+            return a + 1.0
+
+        def definition_func(in_field: DeferredType, out_field: gtscript.Field[float]):
+            with computation(PARALLEL), interval(...):
+                out_field = plus_one(in_field)
+
+        # resolve deferred type on stencil
+        resolve_type(definition_func)
+
+        # parse stencil, which will automatically resolve the deferred type in the lazy_function
+        stencil = parse_definition(
+            definition_func,
+            name=inspect.stack()[0][3],
+            module=self.__class__.__name__,
+        )
+
+        # check that constant was properly inlined
+        assert stencil.computations[0].body.stmts[0].value.rhs.value == 1.0
+
+
 class TestAxisSyntax:
     def test_good_syntax(self):
         def definition_func(in_field: gtscript.Field[float], out_field: gtscript.Field[float]):