1#pragma once
2
3#include <torch/csrc/Exceptions.h>
4#include <torch/csrc/jit/frontend/tracer.h>
5#include <torch/csrc/python_headers.h>
6#include <torch/csrc/utils/object_ptr.h>
7#include <torch/csrc/utils/tensor_numpy.h>
8#include <cstdint>
9#include <limits>
10#include <stdexcept>
11
12// largest integer that can be represented consecutively in a double
13const int64_t DOUBLE_INT_MAX = 9007199254740992;
14
15inline PyObject* THPUtils_packInt32(int32_t value) {
16 return PyLong_FromLong(value);
17}
18
19inline PyObject* THPUtils_packInt64(int64_t value) {
20 return PyLong_FromLongLong(value);
21}
22
23inline PyObject* THPUtils_packUInt32(uint32_t value) {
24 return PyLong_FromUnsignedLong(value);
25}
26
27inline PyObject* THPUtils_packUInt64(uint64_t value) {
28 return PyLong_FromUnsignedLongLong(value);
29}
30
31inline PyObject* THPUtils_packDoubleAsInt(double value) {
32 return PyLong_FromDouble(value);
33}
34
35inline bool THPUtils_checkLong(PyObject* obj) {
36#ifdef USE_NUMPY
37 if (torch::utils::is_numpy_int(obj)) {
38 return true;
39 }
40#endif
41
42 return PyLong_Check(obj) && !PyBool_Check(obj);
43}
44
45inline int32_t THPUtils_unpackInt(PyObject* obj) {
46 // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
47 int overflow;
48 long value = PyLong_AsLongAndOverflow(obj, &overflow);
49 if (value == -1 && PyErr_Occurred()) {
50 throw python_error();
51 }
52 if (overflow != 0) {
53 throw std::runtime_error("Overflow when unpacking long");
54 }
55 if (value > std::numeric_limits<int32_t>::max() ||
56 value < std::numeric_limits<int32_t>::min()) {
57 throw std::runtime_error("Overflow when unpacking long");
58 }
59 return (int32_t)value;
60}
61
62inline int64_t THPUtils_unpackLong(PyObject* obj) {
63 // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
64 int overflow;
65 long long value = PyLong_AsLongLongAndOverflow(obj, &overflow);
66 if (value == -1 && PyErr_Occurred()) {
67 throw python_error();
68 }
69 if (overflow != 0) {
70 throw std::runtime_error("Overflow when unpacking long");
71 }
72 return (int64_t)value;
73}
74
75inline uint32_t THPUtils_unpackUInt32(PyObject* obj) {
76 unsigned long value = PyLong_AsUnsignedLong(obj);
77 if (PyErr_Occurred()) {
78 throw python_error();
79 }
80 if (value > std::numeric_limits<uint32_t>::max()) {
81 throw std::runtime_error("Overflow when unpacking unsigned long");
82 }
83 return (uint32_t)value;
84}
85
86inline uint64_t THPUtils_unpackUInt64(PyObject* obj) {
87 unsigned long long value = PyLong_AsUnsignedLongLong(obj);
88 if (PyErr_Occurred()) {
89 throw python_error();
90 }
91 return (uint64_t)value;
92}
93
94inline bool THPUtils_checkIndex(PyObject* obj) {
95 if (PyBool_Check(obj)) {
96 return false;
97 }
98 if (THPUtils_checkLong(obj)) {
99 return true;
100 }
101 torch::jit::tracer::NoWarn no_warn_guard;
102 auto index = THPObjectPtr(PyNumber_Index(obj));
103 if (!index) {
104 PyErr_Clear();
105 return false;
106 }
107 return true;
108}
109
110inline int64_t THPUtils_unpackIndex(PyObject* obj) {
111 if (!THPUtils_checkLong(obj)) {
112 auto index = THPObjectPtr(PyNumber_Index(obj));
113 if (index == nullptr) {
114 throw python_error();
115 }
116 // NB: This needs to be called before `index` goes out of scope and the
117 // underlying object's refcount is decremented
118 return THPUtils_unpackLong(index.get());
119 }
120 return THPUtils_unpackLong(obj);
121}
122
123inline bool THPUtils_unpackBool(PyObject* obj) {
124 if (obj == Py_True) {
125 return true;
126 } else if (obj == Py_False) {
127 return false;
128 } else {
129 throw std::runtime_error("couldn't convert python object to boolean");
130 }
131}
132
133inline bool THPUtils_checkDouble(PyObject* obj) {
134#ifdef USE_NUMPY
135 if (torch::utils::is_numpy_scalar(obj)) {
136 return true;
137 }
138#endif
139 return PyFloat_Check(obj) || PyLong_Check(obj);
140}
141
142inline double THPUtils_unpackDouble(PyObject* obj) {
143 if (PyFloat_Check(obj)) {
144 return PyFloat_AS_DOUBLE(obj);
145 }
146 double value = PyFloat_AsDouble(obj);
147 if (value == -1 && PyErr_Occurred()) {
148 throw python_error();
149 }
150 return value;
151}
152
153inline c10::complex<double> THPUtils_unpackComplexDouble(PyObject* obj) {
154 Py_complex value = PyComplex_AsCComplex(obj);
155 if (value.real == -1.0 && PyErr_Occurred()) {
156 throw python_error();
157 }
158
159 return c10::complex<double>(value.real, value.imag);
160}
161
162inline bool THPUtils_unpackNumberAsBool(PyObject* obj) {
163 if (PyFloat_Check(obj)) {
164 return (bool)PyFloat_AS_DOUBLE(obj);
165 }
166
167 if (PyComplex_Check(obj)) {
168 double real_val = PyComplex_RealAsDouble(obj);
169 double imag_val = PyComplex_ImagAsDouble(obj);
170 return !(real_val == 0 && imag_val == 0);
171 }
172
173 // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
174 int overflow;
175 long long value = PyLong_AsLongLongAndOverflow(obj, &overflow);
176 if (value == -1 && PyErr_Occurred()) {
177 throw python_error();
178 }
179 // No need to check overflow, because when overflow occured, it should
180 // return true in order to keep the same behavior of numpy.
181 return (bool)value;
182}
183