1from types import ModuleType
2
3import weakref
4
5from numba.core.errors import ConstantInferenceError, NumbaError
6from numba.core import ir
7
8
9class ConstantInference(object):
10    """
11    A constant inference engine for a given interpreter.
12    Inference inspects the IR to try and compute a compile-time constant for
13    a variable.
14
15    This shouldn't be used directly, instead call Interpreter.infer_constant().
16    """
17
18    def __init__(self, func_ir):
19        # Avoid cyclic references as some user-visible objects may be
20        # held alive in the cache
21        self._func_ir = weakref.proxy(func_ir)
22        self._cache = {}
23
24    def infer_constant(self, name, loc=None):
25        """
26        Infer a constant value for the given variable *name*.
27        If no value can be inferred, numba.errors.ConstantInferenceError
28        is raised.
29        """
30        if name not in self._cache:
31            try:
32                self._cache[name] = (True, self._do_infer(name))
33            except ConstantInferenceError as exc:
34                # Store the exception args only, to avoid keeping
35                # a whole traceback alive.
36                self._cache[name] = (False, (exc.__class__, exc.args))
37        success, val = self._cache[name]
38        if success:
39            return val
40        else:
41            exc, args = val
42            if issubclass(exc, NumbaError):
43                raise exc(*args, loc=loc)
44            else:
45                raise exc(*args)
46
47    def _fail(self, val):
48        # The location here is set to None because `val` is the ir.Var name
49        # and not the actual offending use of the var. When this is raised it is
50        # caught in the flow control of `infer_constant` and the class and args
51        # (the message) are captured and then raised again but with the location
52        # set to the expression that caused the constant inference error.
53        raise ConstantInferenceError(
54            "Constant inference not possible for: %s" % (val,), loc=None)
55
56    def _do_infer(self, name):
57        if not isinstance(name, str):
58            raise TypeError("infer_constant() called with non-str %r"
59                            % (name,))
60        try:
61            defn = self._func_ir.get_definition(name)
62        except KeyError:
63            raise ConstantInferenceError(
64                "no single definition for %r" % (name,))
65        try:
66            const = defn.infer_constant()
67        except ConstantInferenceError:
68            if isinstance(defn, ir.Expr):
69                return self._infer_expr(defn)
70            self._fail(defn)
71        return const
72
73    def _infer_expr(self, expr):
74        # Infer an expression: handle supported cases
75        if expr.op == 'call':
76            func = self.infer_constant(expr.func.name, loc=expr.loc)
77            return self._infer_call(func, expr)
78        elif expr.op == 'getattr':
79            value = self.infer_constant(expr.value.name, loc=expr.loc)
80            return self._infer_getattr(value, expr)
81        elif expr.op == 'build_list':
82            return [self.infer_constant(i.name, loc=expr.loc) for i in
83                    expr.items]
84        elif expr.op == 'build_tuple':
85            return tuple(self.infer_constant(i.name, loc=expr.loc) for i in
86                         expr.items)
87        self._fail(expr)
88
89    def _infer_call(self, func, expr):
90        if expr.kws or expr.vararg:
91            self._fail(expr)
92        # Check supported callables
93        _slice = func in (slice,)
94        _exc = isinstance(func, type) and issubclass(func, BaseException)
95        if _slice or _exc:
96            args = [self.infer_constant(a.name, loc=expr.loc) for a in
97                    expr.args]
98            if _slice:
99                return func(*args)
100            elif _exc:
101                # If the exception class is user defined it may implement a ctor
102                # that does not pass the args to the super. Therefore return the
103                # raw class and the args so this can be instantiated at the call
104                # site in the way the user source expects it to be.
105                return func, args
106            else:
107                assert 0, 'Unreachable'
108
109        self._fail(expr)
110
111    def _infer_getattr(self, value, expr):
112        if isinstance(value, (ModuleType, type)):
113            # Allow looking up a constant on a class or module
114            return getattr(value, expr.attr)
115        self._fail(expr)
116