1import cmath
2
3import numpy as np
4
5from numba import float32
6from numba.pycc import CC, exportmany, export
7from numba.tests.support import has_blas
8from numba import typed
9
10
11#
12# New API
13#
14
15cc = CC('pycc_test_simple')
16cc.use_nrt = False
17
18# Note the first signature omits the return type
19@cc.export('multf', (float32, float32))
20@cc.export('multi', 'i4(i4, i4)')
21def mult(a, b):
22    return a * b
23
24# Test imported C globals such as Py_None, PyExc_ZeroDivisionError
25@cc.export('get_none', 'none()')
26def get_none():
27    return None
28
29@cc.export('div', 'f8(f8, f8)')
30def div(x, y):
31    return x / y
32
33_two = 2
34
35# This one can't be compiled by the legacy API as it doesn't execute
36# the script in a proper module.
37@cc.export('square', 'i8(i8)')
38def square(u):
39    return u ** _two
40
41# These ones need helperlib
42cc_helperlib = CC('pycc_test_helperlib')
43cc_helperlib.use_nrt = False
44
45@cc_helperlib.export('power', 'i8(i8, i8)')
46def power(u, v):
47    return u ** v
48
49@cc_helperlib.export('sqrt', 'c16(c16)')
50def sqrt(u):
51    return cmath.sqrt(u)
52
53@cc_helperlib.export('size', 'i8(f8[:])')
54def size(arr):
55    return arr.size
56
57# Exercise linking to Numpy math functions
58@cc_helperlib.export('np_sqrt', 'f8(f8)')
59def np_sqrt(u):
60    return np.sqrt(u)
61
62@cc_helperlib.export('spacing', 'f8(f8)')
63def np_spacing(u):
64    return np.spacing(u)
65
66
67# This one clashes with libc random() unless pycc is careful with naming.
68@cc_helperlib.export('random', 'f8(i4)')
69def random_impl(seed):
70    if seed != -1:
71        np.random.seed(seed)
72    return np.random.random()
73
74# These ones need NRT
75cc_nrt = CC('pycc_test_nrt')
76
77@cc_nrt.export('zero_scalar', 'f8(i4)')
78def zero_scalar(n):
79    arr = np.zeros(n)
80    return arr[-1]
81
82if has_blas:
83    # This one also needs BLAS
84    @cc_nrt.export('vector_dot', 'f8(i4)')
85    def vector_dot(n):
86        a = np.linspace(1, n, n)
87        return np.dot(a, a)
88
89# This one needs an environment
90@cc_nrt.export('zeros', 'f8[:](i4)')
91def zeros(n):
92    return np.zeros(n)
93
94# requires list dtor, #issue3535
95@cc_nrt.export('np_argsort', 'intp[:](float64[:])')
96def np_argsort(arr):
97    return np.argsort(arr)
98
99#
100# Legacy API
101#
102
103exportmany(['multf f4(f4,f4)', 'multi i4(i4,i4)'])(mult)
104# Needs to link to helperlib to due with complex arguments
105# export('multc c16(c16,c16)')(mult)
106export('mult f8(f8, f8)')(mult)
107
108
109@cc_nrt.export('dict_usecase', 'intp[:](intp[:])')
110def dict_usecase(arr):
111    d = typed.Dict()
112    for i in range(arr.size):
113        d[i] = arr[i]
114    out = np.zeros_like(arr)
115    for k, v in d.items():
116        out[k] = k * v
117    return out
118