Repository: GridTools/gt4py
Pull request #1659 (2024-09-23 14:43:38)

feat[next]: unary operator for number type cast
edits in reference to issue #1643 

```python
    def unary_float():
        return float(-1)
```
 as well as other type conversions


Patch:
diff --git a/src/gt4py/next/ffront/foast_to_gtir.py b/src/gt4py/next/ffront/foast_to_gtir.py
index 4d3230a5..2e007c28 100644
--- a/src/gt4py/next/ffront/foast_to_gtir.py
+++ b/src/gt4py/next/ffront/foast_to_gtir.py
@@ -405,16 +405,23 @@ class FieldOperatorLowering(eve.PreserveLocationVisitor, eve.NodeTranslator):
         return self._make_reduction_expr(node, "minimum", init_expr, **kwargs)
 
     def _visit_type_constr(self, node: foast.Call, **kwargs: Any) -> itir.Expr:
-        if isinstance(node.args[0], foast.Constant):
-            node_kind = self.visit(node.type).kind.name.lower()
-            target_type = fbuiltins.BUILTINS[node_kind]
-            source_type = {**fbuiltins.BUILTINS, "string": str}[node.args[0].type.__str__().lower()]
-            if target_type is bool and source_type is not bool:
-                return im.literal(str(bool(source_type(node.args[0].value))), "bool")
-            return im.literal(str(node.args[0].value), node_kind)
-        raise FieldOperatorLoweringError(
-            f"Encountered a type cast, which is not supported: {node}."
-        )
+        el = node.args[0]
+        node_kind = self.visit(node.type).kind.name.lower()
+        source_type = {**fbuiltins.BUILTINS, "string": str}[el.type.__str__().lower()]
+        target_type = fbuiltins.BUILTINS[node_kind]
+
+        if isinstance(el, foast.Constant):
+            val = source_type(el.value)
+        elif isinstance(el, foast.UnaryOp) and isinstance(el.operand, foast.Constant):
+            operand = source_type(el.operand.value)
+            val = eval(f"lambda arg: {el.op}arg")(operand)
+        else:
+            raise FieldOperatorLoweringError(
+                f"Type cast only supports literal arguments, {node.type} not supported."
+            )
+        val = target_type(val)
+
+        return im.literal(str(val), node_kind)
 
     def _make_literal(self, val: Any, type_: ts.TypeSpec) -> itir.Expr:
         if isinstance(type_, ts.TupleType):
diff --git a/src/gt4py/next/ffront/foast_to_itir.py b/src/gt4py/next/ffront/foast_to_itir.py
index b32bf744..7936eda1 100644
--- a/src/gt4py/next/ffront/foast_to_itir.py
+++ b/src/gt4py/next/ffront/foast_to_itir.py
@@ -447,18 +447,23 @@ class FieldOperatorLowering(PreserveLocationVisitor, NodeTranslator):
         return self._make_reduction_expr(node, "minimum", init_expr, **kwargs)
 
     def _visit_type_constr(self, node: foast.Call, **kwargs: Any) -> itir.Expr:
-        if isinstance(node.args[0], foast.Constant):
-            node_kind = self.visit(node.type).kind.name.lower()
-            target_type = fbuiltins.BUILTINS[node_kind]
-            source_type = {**fbuiltins.BUILTINS, "string": str}[node.args[0].type.__str__().lower()]
-            if target_type is bool and source_type is not bool:
-                return im.promote_to_const_iterator(
-                    im.literal(str(bool(source_type(node.args[0].value))), "bool")
-                )
-            return im.promote_to_const_iterator(im.literal(str(node.args[0].value), node_kind))
-        raise FieldOperatorLoweringError(
-            f"Encountered a type cast, which is not supported: {node}."
-        )
+        el = node.args[0]
+        node_kind = self.visit(node.type).kind.name.lower()
+        source_type = {**fbuiltins.BUILTINS, "string": str}[el.type.__str__().lower()]
+        target_type = fbuiltins.BUILTINS[node_kind]
+
+        if isinstance(el, foast.Constant):
+            val = source_type(el.value)
+        elif isinstance(el, foast.UnaryOp) and isinstance(el.operand, foast.Constant):
+            operand = source_type(el.operand.value)
+            val = eval(f"lambda arg: {el.op}arg")(operand)
+        else:
+            raise FieldOperatorLoweringError(
+                f"Type cast only supports literal arguments, {node.type} not supported."
+            )
+        val = target_type(val)
+
+        return im.promote_to_const_iterator(im.literal(str(val), node_kind))
 
     def _make_literal(self, val: Any, type_: ts.TypeSpec) -> itir.Expr:
         # TODO(havogt): lifted nullary lambdas are not supported in iterator.embedded due to an implementation detail;
diff --git a/src/gt4py/next/ffront/func_to_foast.py b/src/gt4py/next/ffront/func_to_foast.py
index 887e6cec..ebe12d3a 100644
--- a/src/gt4py/next/ffront/func_to_foast.py
+++ b/src/gt4py/next/ffront/func_to_foast.py
@@ -503,11 +503,16 @@ class FieldOperatorParser(DialectParser[foast.FunctionDefinition]):
         return foast.CompareOperator.NOTEQ
 
     def _verify_builtin_type_constructor(self, node: ast.Call) -> None:
-        if len(node.args) > 0 and not isinstance(node.args[0], ast.Constant):
-            raise errors.DSLError(
-                self.get_location(node),
-                f"'{self._func_name(node)}()' only takes literal arguments.",
-            )
+        if len(node.args) > 0:
+            arg = node.args[0]
+            if not (
+                isinstance(arg, ast.Constant)
+                or (isinstance(arg, ast.UnaryOp) and isinstance(arg.operand, ast.Constant))
+            ):
+                raise errors.DSLError(
+                    self.get_location(node),
+                    f"'{self._func_name(node)}()' only takes literal arguments.",
+                )
 
     def _func_name(self, node: ast.Call) -> str:
         return node.func.id  # type: ignore[attr-defined] # We want this to fail if the attribute does not exist unexpectedly.


Test patch:
diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_unary_builtins.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_unary_builtins.py
index a6eff8db..89c341e9 100644
--- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_unary_builtins.py
+++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_unary_builtins.py
@@ -11,6 +11,7 @@ import pytest
 
 import gt4py.next as gtx
 from gt4py.next import (
+    broadcast,
     cbrt,
     ceil,
     cos,
@@ -127,6 +128,30 @@ def test_unary_neg(cartesian_case):
     cases.verify_with_default_data(cartesian_case, uneg, ref=lambda inp1: -inp1)
 
 
+def test_unary_neg_float_conversion(cartesian_case):
+    @gtx.field_operator
+    def uneg_float() -> cases.IFloatField:
+        inp_f = broadcast(float(-1), (IDim,))
+        return inp_f
+
+    size = cartesian_case.default_sizes[IDim]
+    ref = cartesian_case.as_field([IDim], np.full(size, -1.0, dtype=float))
+    out = cases.allocate(cartesian_case, uneg_float, cases.RETURN)()
+    cases.verify(cartesian_case, uneg_float, out=out, ref=ref)
+
+
+def test_unary_neg_bool_conversion(cartesian_case):
+    @gtx.field_operator
+    def uneg_bool() -> cases.IBoolField:
+        inp_f = broadcast(bool(-1), (IDim,))
+        return inp_f
+
+    size = cartesian_case.default_sizes[IDim]
+    ref = cartesian_case.as_field([IDim], np.full(size, True, dtype=bool))
+    out = cases.allocate(cartesian_case, uneg_bool, cases.RETURN)()
+    cases.verify(cartesian_case, uneg_bool, out=out, ref=ref)
+
+
 def test_unary_invert(cartesian_case):
     @gtx.field_operator
     def tilde_fieldop(inp1: cases.IBoolField) -> cases.IBoolField:
diff --git a/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py b/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py
index 081c2514..706de8a3 100644
--- a/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py
+++ b/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py
@@ -395,6 +395,22 @@ def test_unary_plus():
     assert lowered.expr == reference
 
 
+@pytest.mark.parametrize("var, var_type", [("-1.0", "float64"), ("True", "bool")])
+def test_unary_op_type_conversion(var, var_type):
+    def unary_float():
+        return float(-1)
+
+    def unary_bool():
+        return bool(-1)
+
+    fun = unary_bool if var_type == "bool" else unary_float
+    parsed = FieldOperatorParser.apply_to_function(fun)
+    lowered = FieldOperatorLowering.apply(parsed)
+    reference = im.literal(var, var_type)
+
+    assert lowered.expr == reference
+
+
 def test_unpacking():
     """Unpacking assigns should get separated."""
 
@@ -862,9 +878,9 @@ def test_builtin_float_constructors():
         im.literal("0.1", "float64"),
         im.literal("0.1", "float32"),
         im.literal("0.1", "float64"),
-        im.literal(".1", "float64"),
-        im.literal(".1", "float32"),
-        im.literal(".1", "float64"),
+        im.literal("0.1", "float64"),
+        im.literal("0.1", "float32"),
+        im.literal("0.1", "float64"),
     )
 
     assert lowered.expr == reference
diff --git a/tests/next_tests/unit_tests/ffront_tests/test_foast_to_itir.py b/tests/next_tests/unit_tests/ffront_tests/test_foast_to_itir.py
index faedcf15..c102df9d 100644
--- a/tests/next_tests/unit_tests/ffront_tests/test_foast_to_itir.py
+++ b/tests/next_tests/unit_tests/ffront_tests/test_foast_to_itir.py
@@ -177,6 +177,22 @@ def test_unary_ops():
     assert lowered.expr == reference
 
 
+@pytest.mark.parametrize("var, var_type", [("-1.0", "float64"), ("True", "bool")])
+def test_unary_op_type_conversion(var, var_type):
+    def unary_float():
+        return float(-1)
+
+    def unary_bool():
+        return bool(-1)
+
+    fun = unary_bool if var_type == "bool" else unary_float
+    parsed = FieldOperatorParser.apply_to_function(fun)
+    lowered = FieldOperatorLowering.apply(parsed)
+    reference = im.promote_to_const_iterator(im.literal(var, var_type))
+
+    assert lowered.expr == reference
+
+
 def test_unpacking():
     """Unpacking assigns should get separated."""
 
@@ -553,9 +569,9 @@ def test_builtin_float_constructors():
         im.promote_to_const_iterator(im.literal("0.1", "float64")),
         im.promote_to_const_iterator(im.literal("0.1", "float32")),
         im.promote_to_const_iterator(im.literal("0.1", "float64")),
-        im.promote_to_const_iterator(im.literal(".1", "float64")),
-        im.promote_to_const_iterator(im.literal(".1", "float32")),
-        im.promote_to_const_iterator(im.literal(".1", "float64")),
+        im.promote_to_const_iterator(im.literal("0.1", "float64")),
+        im.promote_to_const_iterator(im.literal("0.1", "float32")),
+        im.promote_to_const_iterator(im.literal("0.1", "float64")),
     )
 
     assert lowered.expr == reference