From f76edf725ea1125ed0bbcb2799c68a788b518924 Mon Sep 17 00:00:00 2001
From: Dominik Thoennes <dominik.thoennes@fau.de>
Date: Fri, 8 Jan 2021 14:04:16 +0100
Subject: [PATCH] add single argument integer functions
---
pystencils/integer_functions.py | 25 +++++++++++++++++++++++--
1 file changed, 23 insertions(+), 2 deletions(-)
diff --git a/pystencils/integer_functions.py b/pystencils/integer_functions.py
index a6b9daf21..a0516ac3a 100644
--- a/pystencils/integer_functions.py
+++ b/pystencils/integer_functions.py
@@ -27,6 +27,27 @@ class IntegerFunctionTwoArgsMixIn(sp.Function):
raise ValueError("Integer functions can only be constructed with typed expressions")
return super().__new__(cls, *args)
+class IntegerFunctionOneArgMixIn(sp.Function):
+ is_integer = True
+
+ def __new__(cls, arg1):
+ args = []
+
+ if isinstance(arg1, sp.Number) or isinstance(arg1, int):
+ args.append(cast_func(arg1, create_type("int")))
+ elif isinstance(arg1, np.generic):
+ args.append(cast_func(arg1, arg1.dtype))
+ else:
+ args.append(arg1)
+
+
+ try:
+ type = get_type_of_expression(arg1)
+ if not type.is_int():
+ raise ValueError("Argument to integer function is not an int but " + str(type))
+ except NotImplementedError:
+ raise ValueError("Integer functions can only be constructed with typed expressions")
+ return super().__new__(cls, *args)
# noinspection PyPep8Naming
class bitwise_xor(IntegerFunctionTwoArgsMixIn):
@@ -59,11 +80,11 @@ class int_div(IntegerFunctionTwoArgsMixIn):
# noinspection PyPep8Naming
-class int_power_of_2(sp.Function):
+class int_power_of_2(IntegerFunctionOneArgMixIn):
pass
# noinspection PyPep8Naming
-class post_increment(sp.Function):
+class post_increment(IntegerFunctionOneArgMixIn):
pass
--
GitLab