Fix Numba

data_tests/saved__backend__/numba/add_inline.py

@@ -16,8 +16,5 @@ def use_add(n=10000):
     return tmp
 
 
-# __protected__ @njit(cache=True, fastmath=True)
-
-
 def __transonic__():
     return "0.7.1"

data_tests/saved__backend__/numba/assign_func_boost.py

@@ -6,8 +6,5 @@ def func(x):
     return x**2
 
 
-# __protected__ @njit(cache=True, fastmath=True)
-
-
 def __transonic__():
     return "0.7.1"

data_tests/saved__backend__/numba/block_fluidsim.py

@@ -12,9 +12,6 @@ def rk2_step0(state_spect_n12, state_spect, tendencies_n, diss2, dt):
     state_spect_n12[:] = (state_spect + dt / 2 * tendencies_n) * diss2
 
 
-# __protected__ @njit(cache=True, fastmath=True)
-
-
 def arguments_blocks():
     return {
         "rk2_step0": [
@@ -27,8 +24,5 @@ def arguments_blocks():
     }
 
 
-# __protected__ @njit(cache=True, fastmath=True)
-
-
 def __transonic__():
     return "0.7.1"

data_tests/saved__backend__/numba/blocks_type_hints.py

@@ -15,15 +15,9 @@ def block0(a, b, n):
     return result
 
 
-# __protected__ @njit(cache=True, fastmath=True)
-
-
 def arguments_blocks():
     return {"block0": ["a", "b", "n"]}
 
 
-# __protected__ @njit(cache=True, fastmath=True)
-
-
 def __transonic__():
     return "0.7.1"

data_tests/saved__backend__/numba/boosted_class_use_import.py

@@ -9,15 +9,9 @@ def __for_method__MyClass2__myfunc(self_attr0, self_attr1, arg):
     return self_attr1 + self_attr0 + np.abs(arg) + func_import()
 
 
-# __protected__ @njit(cache=True, fastmath=True)
-
-
 def __code_new_method__MyClass2__myfunc():
     return "\n\ndef new_method(self, arg):\n    return backend_func(self.attr0, self.attr1, arg)\n\n"
 
 
-# __protected__ @njit(cache=True, fastmath=True)
-
-
 def __transonic__():
     return "0.7.1"

data_tests/saved__backend__/numba/boosted_func_use_import.py

@@ -9,8 +9,5 @@ def func(a, b):
     return (a * np.log(b)).max() + func_import()
 
 
-# __protected__ @njit(cache=True, fastmath=True)
-
-
 def __transonic__():
     return "0.7.1"

data_tests/saved__backend__/numba/class_blocks.py

@@ -40,15 +40,9 @@ def block1(a, b, n):
     return result
 
 
-# __protected__ @njit(cache=True, fastmath=True)
-
-
 def arguments_blocks():
     return {"block0": ["a", "b", "n"], "block1": ["a", "b", "n"]}
 
 
-# __protected__ @njit(cache=True, fastmath=True)
-
-
 def __transonic__():
     return "0.7.1"

data_tests/saved__backend__/numba/class_rec_calls.py

@@ -16,15 +16,9 @@ def __for_method__Myclass__func(self_attr, self_attr2, arg):
         )
 
 
-# __protected__ @njit(cache=True, fastmath=True)
-
-
 def __code_new_method__Myclass__func():
     return "\n\ndef new_method(self, arg):\n    return backend_func(self.attr, self.attr2, arg)\n\n"
 
 
-# __protected__ @njit(cache=True, fastmath=True)
-
-
 def __transonic__():
     return "0.7.1"

data_tests/saved__backend__/numba/classic.py

@@ -8,8 +8,5 @@ def func(a, b):
     return (a * np.log(b)).max()
 
 
-# __protected__ @njit(cache=True, fastmath=True)
-
-
 def __transonic__():
     return "0.7.1"

data_tests/saved__backend__/numba/default_params.py

@@ -7,8 +7,5 @@ def func(a=1, b=None, c=1.0):
     return a + c
 
 
-# __protected__ @njit(cache=True, fastmath=True)
-
-
 def __transonic__():
     return "0.7.1"

data_tests/saved__backend__/numba/methods.py

@@ -9,15 +9,9 @@ def __for_method__Transmitter____call__(self_arr, self_freq, inp):
     return (inp * np.exp(np.arange(len(inp)) * self_freq * 1j), self_arr)
 
 
-# __protected__ @njit(cache=True, fastmath=True)
-
-
 def __code_new_method__Transmitter____call__():
     return "\n\ndef new_method(self, inp):\n    return backend_func(self.arr, self.freq, inp)\n\n"
 
 
-# __protected__ @njit(cache=True, fastmath=True)
-
-
 def __transonic__():
     return "0.7.1"

data_tests/saved__backend__/numba/mixed_classic_type_hint.py

@@ -15,8 +15,5 @@ def func1(a, b):
     return a * np.cos(b)
 
 
-# __protected__ @njit(cache=True, fastmath=True)
-
-
 def __transonic__():
     return "0.7.1"

data_tests/saved__backend__/numba/no_arg.py

@@ -13,8 +13,5 @@ def func2():
     return 1
 
 
-# __protected__ @njit(cache=True, fastmath=True)
-
-
 def __transonic__():
     return "0.7.1"

data_tests/saved__backend__/numba/row_sum_boost.py

@@ -24,8 +24,5 @@ def row_sum_loops(arr, columns):
     return res
 
 
-# __protected__ @njit(cache=True, fastmath=True)
-
-
 def __transonic__():
     return "0.7.1"

data_tests/saved__backend__/numba/subpackages.py

@@ -30,8 +30,5 @@ def test_sp_special(v, x):
     return jv(v, x)
 
 
-# __protected__ @njit(cache=True, fastmath=True)
-
-
 def __transonic__():
     return "0.7.1"

data_tests/saved__backend__/numba/type_hint_notemplate.py

@@ -10,8 +10,5 @@ def compute(a, b, c, d, e):
     return tmp
 
 
-# __protected__ @njit(cache=True, fastmath=True)
-
-
 def __transonic__():
     return "0.7.1"

src/transonic/aheadoftime.py

@@ -62,9 +62,9 @@ def _get_data_from_func_or_data(data_or_func):
     which leads to issue with Pythran on Windows.
     They are now saved in functions, but we continue supporting the "old" extensions.
     """
-    try:
+    if callable(data_or_func):
         return data_or_func()
-    except TypeError:
+    else:
         # extension compiled with Transonic < 0.7.2
         return data_or_func
 

src/transonic/backends/numba.py

@@ -28,7 +28,15 @@ def add_numba_comments(code):
     new_body = [CommentLine("# __protected__ from numba import njit")]
 
     for node in mod.body:
-        if isinstance(node, gast.FunctionDef):
+        if (
+            isinstance(node, gast.FunctionDef)
+            and node.name
+            not in (
+                "arguments_blocks",
+                "__transonic__",
+            )
+            and not node.name.startswith("__code_new_method__")
+        ):
             new_body.append(
                 CommentLine("# __protected__ @njit(cache=True, fastmath=True)")
             )