py_dict.cpp 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683
  1. #include "py_dict.h"
  2. #include "py_iterator.h"
  3. #include "py_cast.h"
  4. #include "py_errors.h"
  5. #include "py_utils.h"
  6. #include <yql/essentials/public/udf/udf_value.h>
  7. #include <yql/essentials/public/udf/udf_value_builder.h>
  8. #include <yql/essentials/public/udf/udf_type_inspection.h>
  9. using namespace NKikimr;
  10. namespace NPython {
  11. //////////////////////////////////////////////////////////////////////////////
  12. // TPyLazyDict interface
  13. //////////////////////////////////////////////////////////////////////////////
  14. struct TPyLazyDict
  15. {
  16. using TPtr = NUdf::TRefCountedPtr<TPyLazyDict, TPyPtrOps<TPyLazyDict>>;
  17. PyObject_HEAD;
  18. TPyCastContext::TPtr CastCtx;
  19. const NUdf::TType* KeyType;
  20. const NUdf::TType* PayloadType;
  21. TPyCleanupListItem<NUdf::IBoxedValuePtr> Value;
  22. inline static TPyLazyDict* Cast(PyObject* o) {
  23. return reinterpret_cast<TPyLazyDict*>(o);
  24. }
  25. inline static void Dealloc(PyObject* self) {
  26. delete Cast(self);
  27. }
  28. static PyObject* New(
  29. const TPyCastContext::TPtr& castCtx,
  30. const NUdf::TType* keyType,
  31. const NUdf::TType* payloadType,
  32. NUdf::IBoxedValuePtr&& value);
  33. static int Bool(PyObject* self);
  34. static PyObject* Repr(PyObject* self);
  35. static Py_ssize_t Len(PyObject* self);
  36. static PyObject* Subscript(PyObject* self, PyObject* key);
  37. static int Contains(PyObject* self, PyObject* key);
  38. static PyObject* Get(PyObject* self, PyObject* args);
  39. static PyObject* Iter(PyObject* self) { return Keys(self, nullptr); }
  40. static PyObject* Keys(PyObject* self, PyObject* /* args */);
  41. static PyObject* Items(PyObject* self, PyObject* /* args */);
  42. static PyObject* Values(PyObject* self, PyObject* /* args */);
  43. };
  44. PyMappingMethods LazyDictMapping = {
  45. INIT_MEMBER(mp_length, TPyLazyDict::Len),
  46. INIT_MEMBER(mp_subscript, TPyLazyDict::Subscript),
  47. INIT_MEMBER(mp_ass_subscript, nullptr),
  48. };
  49. PySequenceMethods LazyDictSequence = {
  50. INIT_MEMBER(sq_length , TPyLazyDict::Len),
  51. INIT_MEMBER(sq_concat , nullptr),
  52. INIT_MEMBER(sq_repeat , nullptr),
  53. INIT_MEMBER(sq_item , nullptr),
  54. #if PY_MAJOR_VERSION >= 3
  55. INIT_MEMBER(was_sq_slice , nullptr),
  56. #else
  57. INIT_MEMBER(sq_slice , nullptr),
  58. #endif
  59. INIT_MEMBER(sq_ass_item , nullptr),
  60. #if PY_MAJOR_VERSION >= 3
  61. INIT_MEMBER(was_sq_ass_slice , nullptr),
  62. #else
  63. INIT_MEMBER(sq_ass_slice , nullptr),
  64. #endif
  65. INIT_MEMBER(sq_contains , TPyLazyDict::Contains),
  66. INIT_MEMBER(sq_inplace_concat , nullptr),
  67. INIT_MEMBER(sq_inplace_repeat , nullptr),
  68. };
  69. PyNumberMethods LazyDictNumbering = {
  70. INIT_MEMBER(nb_add, nullptr),
  71. INIT_MEMBER(nb_subtract, nullptr),
  72. INIT_MEMBER(nb_multiply, nullptr),
  73. #if PY_MAJOR_VERSION < 3
  74. INIT_MEMBER(nb_divide, nullptr),
  75. #endif
  76. INIT_MEMBER(nb_remainder, nullptr),
  77. INIT_MEMBER(nb_divmod, nullptr),
  78. INIT_MEMBER(nb_power, nullptr),
  79. INIT_MEMBER(nb_negative, nullptr),
  80. INIT_MEMBER(nb_positive, nullptr),
  81. INIT_MEMBER(nb_absolute, nullptr),
  82. #if PY_MAJOR_VERSION >= 3
  83. INIT_MEMBER(nb_bool, TPyLazyDict::Bool),
  84. #else
  85. INIT_MEMBER(nb_nonzero, TPyLazyDict::Bool),
  86. #endif
  87. INIT_MEMBER(nb_invert, nullptr),
  88. INIT_MEMBER(nb_lshift, nullptr),
  89. INIT_MEMBER(nb_rshift, nullptr),
  90. INIT_MEMBER(nb_and, nullptr),
  91. INIT_MEMBER(nb_xor, nullptr),
  92. INIT_MEMBER(nb_or, nullptr),
  93. #if PY_MAJOR_VERSION < 3
  94. INIT_MEMBER(nb_coerce, nullptr),
  95. #endif
  96. INIT_MEMBER(nb_int, nullptr),
  97. #if PY_MAJOR_VERSION >= 3
  98. INIT_MEMBER(nb_reserved, nullptr),
  99. #else
  100. INIT_MEMBER(nb_long, nullptr),
  101. #endif
  102. INIT_MEMBER(nb_float, nullptr),
  103. #if PY_MAJOR_VERSION < 3
  104. INIT_MEMBER(nb_oct, nullptr),
  105. INIT_MEMBER(nb_hex, nullptr),
  106. #endif
  107. INIT_MEMBER(nb_inplace_add, nullptr),
  108. INIT_MEMBER(nb_inplace_subtract, nullptr),
  109. INIT_MEMBER(nb_inplace_multiply, nullptr),
  110. INIT_MEMBER(nb_inplace_remainder, nullptr),
  111. INIT_MEMBER(nb_inplace_power, nullptr),
  112. INIT_MEMBER(nb_inplace_lshift, nullptr),
  113. INIT_MEMBER(nb_inplace_rshift, nullptr),
  114. INIT_MEMBER(nb_inplace_and, nullptr),
  115. INIT_MEMBER(nb_inplace_xor, nullptr),
  116. INIT_MEMBER(nb_inplace_or, nullptr),
  117. INIT_MEMBER(nb_floor_divide, nullptr),
  118. INIT_MEMBER(nb_true_divide, nullptr),
  119. INIT_MEMBER(nb_inplace_floor_divide, nullptr),
  120. INIT_MEMBER(nb_inplace_true_divide, nullptr),
  121. INIT_MEMBER(nb_index, nullptr),
  122. #if PY_MAJOR_VERSION >= 3
  123. INIT_MEMBER(nb_matrix_multiply, nullptr),
  124. INIT_MEMBER(nb_inplace_matrix_multiply, nullptr),
  125. #endif
  126. };
  127. #if PY_MAJOR_VERSION >= 3
  128. #define Py_TPFLAGS_HAVE_ITER 0
  129. #define Py_TPFLAGS_HAVE_SEQUENCE_IN 0
  130. #endif
  131. PyDoc_STRVAR(get__doc__,
  132. "D.get(k[,d]) -> D[k] if k in D, else d. d defaults to None.");
  133. PyDoc_STRVAR(keys__doc__,
  134. "D.keys() -> an iterator over the keys of D");
  135. PyDoc_STRVAR(values__doc__,
  136. "D.values() -> an iterator over the values of D");
  137. PyDoc_STRVAR(items__doc__,
  138. "D.items() -> an iterator over the (key, value) items of D");
  139. #if PY_MAJOR_VERSION < 3
  140. PyDoc_STRVAR(iterkeys__doc__,
  141. "D.iterkeys() -> an iterator over the keys of D");
  142. PyDoc_STRVAR(itervalues__doc__,
  143. "D.itervalues() -> an iterator over the values of D");
  144. PyDoc_STRVAR(iteritems__doc__,
  145. "D.iteritems() -> an iterator over the (key, value) items of D");
  146. #endif
  147. static PyMethodDef LazyDictMethods[] = {
  148. { "get", TPyLazyDict::Get, METH_VARARGS, get__doc__ },
  149. { "keys", TPyLazyDict::Keys, METH_NOARGS, keys__doc__ },
  150. { "items", TPyLazyDict::Items, METH_NOARGS, items__doc__ },
  151. { "values", TPyLazyDict::Values, METH_NOARGS, values__doc__ },
  152. #if PY_MAJOR_VERSION < 3
  153. { "iterkeys", TPyLazyDict::Keys, METH_NOARGS, iterkeys__doc__ },
  154. { "iteritems", TPyLazyDict::Items, METH_NOARGS, iteritems__doc__ },
  155. { "itervalues", TPyLazyDict::Values, METH_NOARGS, itervalues__doc__ },
  156. #endif
  157. { nullptr, nullptr, 0, nullptr } /* sentinel */
  158. };
  159. PyTypeObject PyLazyDictType = {
  160. PyVarObject_HEAD_INIT(&PyType_Type, 0)
  161. INIT_MEMBER(tp_name , "yql.TDict"),
  162. INIT_MEMBER(tp_basicsize , sizeof(TPyLazyDict)),
  163. INIT_MEMBER(tp_itemsize , 0),
  164. INIT_MEMBER(tp_dealloc , TPyLazyDict::Dealloc),
  165. #if PY_VERSION_HEX < 0x030800b4
  166. INIT_MEMBER(tp_print , nullptr),
  167. #else
  168. INIT_MEMBER(tp_vectorcall_offset, 0),
  169. #endif
  170. INIT_MEMBER(tp_getattr , nullptr),
  171. INIT_MEMBER(tp_setattr , nullptr),
  172. #if PY_MAJOR_VERSION >= 3
  173. INIT_MEMBER(tp_as_async , nullptr),
  174. #else
  175. INIT_MEMBER(tp_compare , nullptr),
  176. #endif
  177. INIT_MEMBER(tp_repr , TPyLazyDict::Repr),
  178. INIT_MEMBER(tp_as_number , &LazyDictNumbering),
  179. INIT_MEMBER(tp_as_sequence , &LazyDictSequence),
  180. INIT_MEMBER(tp_as_mapping , &LazyDictMapping),
  181. INIT_MEMBER(tp_hash , nullptr),
  182. INIT_MEMBER(tp_call , nullptr),
  183. INIT_MEMBER(tp_str , nullptr),
  184. INIT_MEMBER(tp_getattro , nullptr),
  185. INIT_MEMBER(tp_setattro , nullptr),
  186. INIT_MEMBER(tp_as_buffer , nullptr),
  187. INIT_MEMBER(tp_flags , Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_ITER | Py_TPFLAGS_HAVE_SEQUENCE_IN),
  188. INIT_MEMBER(tp_doc , "yql.TDict object"),
  189. INIT_MEMBER(tp_traverse , nullptr),
  190. INIT_MEMBER(tp_clear , nullptr),
  191. INIT_MEMBER(tp_richcompare , nullptr),
  192. INIT_MEMBER(tp_weaklistoffset , 0),
  193. INIT_MEMBER(tp_iter , &TPyLazyDict::Iter),
  194. INIT_MEMBER(tp_iternext , nullptr),
  195. INIT_MEMBER(tp_methods , LazyDictMethods),
  196. INIT_MEMBER(tp_members , nullptr),
  197. INIT_MEMBER(tp_getset , nullptr),
  198. INIT_MEMBER(tp_base , nullptr),
  199. INIT_MEMBER(tp_dict , nullptr),
  200. INIT_MEMBER(tp_descr_get , nullptr),
  201. INIT_MEMBER(tp_descr_set , nullptr),
  202. INIT_MEMBER(tp_dictoffset , 0),
  203. INIT_MEMBER(tp_init , nullptr),
  204. INIT_MEMBER(tp_alloc , nullptr),
  205. INIT_MEMBER(tp_new , nullptr),
  206. INIT_MEMBER(tp_free , nullptr),
  207. INIT_MEMBER(tp_is_gc , nullptr),
  208. INIT_MEMBER(tp_bases , nullptr),
  209. INIT_MEMBER(tp_mro , nullptr),
  210. INIT_MEMBER(tp_cache , nullptr),
  211. INIT_MEMBER(tp_subclasses , nullptr),
  212. INIT_MEMBER(tp_weaklist , nullptr),
  213. INIT_MEMBER(tp_del , nullptr),
  214. INIT_MEMBER(tp_version_tag , 0),
  215. #if PY_MAJOR_VERSION >= 3
  216. INIT_MEMBER(tp_finalize , nullptr),
  217. #endif
  218. #if PY_VERSION_HEX >= 0x030800b1
  219. INIT_MEMBER(tp_vectorcall , nullptr),
  220. #endif
  221. #if PY_VERSION_HEX >= 0x030800b4 && PY_VERSION_HEX < 0x03090000
  222. INIT_MEMBER(tp_print , nullptr),
  223. #endif
  224. };
  225. //////////////////////////////////////////////////////////////////////////////
  226. // TPyLazySet interface
  227. //////////////////////////////////////////////////////////////////////////////
  228. struct TPyLazySet
  229. {
  230. using TPtr = NUdf::TRefCountedPtr<TPyLazySet, TPyPtrOps<TPyLazySet>>;
  231. PyObject_HEAD;
  232. TPyCastContext::TPtr CastCtx;
  233. const NUdf::TType* ItemType;
  234. TPyCleanupListItem<NUdf::IBoxedValuePtr> Value;
  235. inline static TPyLazySet* Cast(PyObject* o) {
  236. return reinterpret_cast<TPyLazySet*>(o);
  237. }
  238. inline static void Dealloc(PyObject* self) {
  239. delete Cast(self);
  240. }
  241. static PyObject* New(
  242. const TPyCastContext::TPtr& castCtx,
  243. const NUdf::TType* itemType,
  244. NUdf::IBoxedValuePtr&& value);
  245. static int Bool(PyObject* self);
  246. static PyObject* Repr(PyObject* self);
  247. static Py_ssize_t Len(PyObject* self);
  248. static int Contains(PyObject* self, PyObject* key);
  249. static PyObject* Get(PyObject* self, PyObject* args);
  250. static PyObject* Iter(PyObject* self);
  251. };
  252. PySequenceMethods LazySetSequence = {
  253. INIT_MEMBER(sq_length , TPyLazySet::Len),
  254. INIT_MEMBER(sq_concat , nullptr),
  255. INIT_MEMBER(sq_repeat , nullptr),
  256. INIT_MEMBER(sq_item , nullptr),
  257. #if PY_MAJOR_VERSION >= 3
  258. INIT_MEMBER(was_sq_slice , nullptr),
  259. #else
  260. INIT_MEMBER(sq_slice , nullptr),
  261. #endif
  262. INIT_MEMBER(sq_ass_item , nullptr),
  263. #if PY_MAJOR_VERSION >= 3
  264. INIT_MEMBER(was_sq_ass_slice , nullptr),
  265. #else
  266. INIT_MEMBER(sq_ass_slice , nullptr),
  267. #endif
  268. INIT_MEMBER(sq_contains , TPyLazySet::Contains),
  269. INIT_MEMBER(sq_inplace_concat , nullptr),
  270. INIT_MEMBER(sq_inplace_repeat , nullptr),
  271. };
  272. PyNumberMethods LazySetNumbering = {
  273. INIT_MEMBER(nb_add, nullptr),
  274. INIT_MEMBER(nb_subtract, nullptr),
  275. INIT_MEMBER(nb_multiply, nullptr),
  276. #if PY_MAJOR_VERSION < 3
  277. INIT_MEMBER(nb_divide, nullptr),
  278. #endif
  279. INIT_MEMBER(nb_remainder, nullptr),
  280. INIT_MEMBER(nb_divmod, nullptr),
  281. INIT_MEMBER(nb_power, nullptr),
  282. INIT_MEMBER(nb_negative, nullptr),
  283. INIT_MEMBER(nb_positive, nullptr),
  284. INIT_MEMBER(nb_absolute, nullptr),
  285. #if PY_MAJOR_VERSION >= 3
  286. INIT_MEMBER(nb_bool, TPyLazySet::Bool),
  287. #else
  288. INIT_MEMBER(nb_nonzero, TPyLazySet::Bool),
  289. #endif
  290. INIT_MEMBER(nb_invert, nullptr),
  291. INIT_MEMBER(nb_lshift, nullptr),
  292. INIT_MEMBER(nb_rshift, nullptr),
  293. INIT_MEMBER(nb_and, nullptr),
  294. INIT_MEMBER(nb_xor, nullptr),
  295. INIT_MEMBER(nb_or, nullptr),
  296. #if PY_MAJOR_VERSION < 3
  297. INIT_MEMBER(nb_coerce, nullptr),
  298. #endif
  299. INIT_MEMBER(nb_int, nullptr),
  300. #if PY_MAJOR_VERSION >= 3
  301. INIT_MEMBER(nb_reserved, nullptr),
  302. #else
  303. INIT_MEMBER(nb_long, nullptr),
  304. #endif
  305. INIT_MEMBER(nb_float, nullptr),
  306. #if PY_MAJOR_VERSION < 3
  307. INIT_MEMBER(nb_oct, nullptr),
  308. INIT_MEMBER(nb_hex, nullptr),
  309. #endif
  310. INIT_MEMBER(nb_inplace_add, nullptr),
  311. INIT_MEMBER(nb_inplace_subtract, nullptr),
  312. INIT_MEMBER(nb_inplace_multiply, nullptr),
  313. INIT_MEMBER(nb_inplace_remainder, nullptr),
  314. INIT_MEMBER(nb_inplace_power, nullptr),
  315. INIT_MEMBER(nb_inplace_lshift, nullptr),
  316. INIT_MEMBER(nb_inplace_rshift, nullptr),
  317. INIT_MEMBER(nb_inplace_and, nullptr),
  318. INIT_MEMBER(nb_inplace_xor, nullptr),
  319. INIT_MEMBER(nb_inplace_or, nullptr),
  320. INIT_MEMBER(nb_floor_divide, nullptr),
  321. INIT_MEMBER(nb_true_divide, nullptr),
  322. INIT_MEMBER(nb_inplace_floor_divide, nullptr),
  323. INIT_MEMBER(nb_inplace_true_divide, nullptr),
  324. INIT_MEMBER(nb_index, nullptr),
  325. #if PY_MAJOR_VERSION >= 3
  326. INIT_MEMBER(nb_matrix_multiply, nullptr),
  327. INIT_MEMBER(nb_inplace_matrix_multiply, nullptr),
  328. #endif
  329. };
  330. PyTypeObject PyLazySetType = {
  331. PyVarObject_HEAD_INIT(&PyType_Type, 0)
  332. INIT_MEMBER(tp_name , "yql.TSet"),
  333. INIT_MEMBER(tp_basicsize , sizeof(TPyLazySet)),
  334. INIT_MEMBER(tp_itemsize , 0),
  335. INIT_MEMBER(tp_dealloc , TPyLazySet::Dealloc),
  336. #if PY_VERSION_HEX < 0x030800b4
  337. INIT_MEMBER(tp_print , nullptr),
  338. #else
  339. INIT_MEMBER(tp_vectorcall_offset, 0),
  340. #endif
  341. INIT_MEMBER(tp_getattr , nullptr),
  342. INIT_MEMBER(tp_setattr , nullptr),
  343. #if PY_MAJOR_VERSION >= 3
  344. INIT_MEMBER(tp_as_async , nullptr),
  345. #else
  346. INIT_MEMBER(tp_compare , nullptr),
  347. #endif
  348. INIT_MEMBER(tp_repr , TPyLazySet::Repr),
  349. INIT_MEMBER(tp_as_number , &LazySetNumbering),
  350. INIT_MEMBER(tp_as_sequence , &LazySetSequence),
  351. INIT_MEMBER(tp_as_mapping , nullptr),
  352. INIT_MEMBER(tp_hash , nullptr),
  353. INIT_MEMBER(tp_call , nullptr),
  354. INIT_MEMBER(tp_str , nullptr),
  355. INIT_MEMBER(tp_getattro , nullptr),
  356. INIT_MEMBER(tp_setattro , nullptr),
  357. INIT_MEMBER(tp_as_buffer , nullptr),
  358. INIT_MEMBER(tp_flags , Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_ITER | Py_TPFLAGS_HAVE_SEQUENCE_IN),
  359. INIT_MEMBER(tp_doc , "yql.TSet object"),
  360. INIT_MEMBER(tp_traverse , nullptr),
  361. INIT_MEMBER(tp_clear , nullptr),
  362. INIT_MEMBER(tp_richcompare , nullptr),
  363. INIT_MEMBER(tp_weaklistoffset , 0),
  364. INIT_MEMBER(tp_iter , &TPyLazySet::Iter),
  365. INIT_MEMBER(tp_iternext , nullptr),
  366. INIT_MEMBER(tp_methods , nullptr),
  367. INIT_MEMBER(tp_members , nullptr),
  368. INIT_MEMBER(tp_getset , nullptr),
  369. INIT_MEMBER(tp_base , nullptr),
  370. INIT_MEMBER(tp_dict , nullptr),
  371. INIT_MEMBER(tp_descr_get , nullptr),
  372. INIT_MEMBER(tp_descr_set , nullptr),
  373. INIT_MEMBER(tp_dictoffset , 0),
  374. INIT_MEMBER(tp_init , nullptr),
  375. INIT_MEMBER(tp_alloc , nullptr),
  376. INIT_MEMBER(tp_new , nullptr),
  377. INIT_MEMBER(tp_free , nullptr),
  378. INIT_MEMBER(tp_is_gc , nullptr),
  379. INIT_MEMBER(tp_bases , nullptr),
  380. INIT_MEMBER(tp_mro , nullptr),
  381. INIT_MEMBER(tp_cache , nullptr),
  382. INIT_MEMBER(tp_subclasses , nullptr),
  383. INIT_MEMBER(tp_weaklist , nullptr),
  384. INIT_MEMBER(tp_del , nullptr),
  385. INIT_MEMBER(tp_version_tag , 0),
  386. #if PY_MAJOR_VERSION >= 3
  387. INIT_MEMBER(tp_finalize , nullptr),
  388. #endif
  389. #if PY_VERSION_HEX >= 0x030800b1
  390. INIT_MEMBER(tp_vectorcall , nullptr),
  391. #endif
  392. #if PY_VERSION_HEX >= 0x030800b4 && PY_VERSION_HEX < 0x03090000
  393. INIT_MEMBER(tp_print , nullptr),
  394. #endif
  395. };
  396. //////////////////////////////////////////////////////////////////////////////
  397. // TPyLazyDict implementation
  398. //////////////////////////////////////////////////////////////////////////////
  399. int TPyLazyDict::Bool(PyObject* self)
  400. {
  401. PY_TRY {
  402. return NUdf::TBoxedValueAccessor::HasDictItems(*Cast(self)->Value.Get()) ? 1 : 0;
  403. } PY_CATCH(-1)
  404. }
  405. PyObject* TPyLazyDict::Repr(PyObject*)
  406. {
  407. return PyRepr("<yql.TDict>").Release();
  408. }
  409. Py_ssize_t TPyLazyDict::Len(PyObject* self)
  410. {
  411. PY_TRY {
  412. return static_cast<Py_ssize_t>(NUdf::TBoxedValueAccessor::GetDictLength(*Cast(self)->Value.Get()));
  413. } PY_CATCH(-1)
  414. }
  415. PyObject* TPyLazyDict::Subscript(PyObject* self, PyObject* key)
  416. {
  417. PY_TRY {
  418. TPyLazyDict* dict = Cast(self);
  419. if (dict->KeyType) {
  420. const auto mkqlKey = FromPyObject(dict->CastCtx, dict->KeyType, key);
  421. if (auto value = NUdf::TBoxedValueAccessor::Lookup(*dict->Value.Get(), mkqlKey)) {
  422. return ToPyObject(dict->CastCtx, dict->PayloadType, value.Release().GetOptionalValue()).Release();
  423. }
  424. const TPyObjectPtr repr = PyObject_Repr(key);
  425. PyErr_SetObject(PyExc_KeyError, repr.Get());
  426. return nullptr;
  427. } else {
  428. if (!PyIndex_Check(key)) {
  429. const TPyObjectPtr type = PyObject_Type(key);
  430. const TPyObjectPtr repr = PyObject_Repr(type.Get());
  431. const TPyObjectPtr error = PyUnicode_FromFormat("Unsupported index object type: %R", repr.Get());
  432. PyErr_SetObject(PyExc_TypeError, error.Get());
  433. return nullptr;
  434. }
  435. const Py_ssize_t index = PyNumber_AsSsize_t(key, PyExc_IndexError);
  436. if (index < 0) {
  437. return nullptr;
  438. }
  439. if (auto value = NUdf::TBoxedValueAccessor::Lookup(*dict->Value.Get(), NUdf::TUnboxedValuePod(ui64(index)))) {
  440. return ToPyObject(dict->CastCtx, dict->PayloadType, value.Release().GetOptionalValue()).Release();
  441. }
  442. const TPyObjectPtr repr = PyObject_Repr(key);
  443. PyErr_SetObject(PyExc_IndexError, repr.Get());
  444. return nullptr;
  445. }
  446. } PY_CATCH(nullptr)
  447. }
  448. // -1 error
  449. // 0 not found
  450. // 1 found
  451. int TPyLazyDict::Contains(PyObject* self, PyObject* key)
  452. {
  453. PY_TRY {
  454. TPyLazyDict* dict = Cast(self);
  455. NUdf::TUnboxedValue mkqlKey;
  456. if (dict->KeyType) {
  457. mkqlKey = FromPyObject(dict->CastCtx, dict->KeyType, key);
  458. } else {
  459. if (!PyIndex_Check(key)) {
  460. const TPyObjectPtr type = PyObject_Type(key);
  461. const TPyObjectPtr repr = PyObject_Repr(type.Get());
  462. const TPyObjectPtr error = PyUnicode_FromFormat("Unsupported index object type: %R", repr.Get());
  463. PyErr_SetObject(PyExc_TypeError, error.Get());
  464. return -1;
  465. }
  466. const Py_ssize_t index = PyNumber_AsSsize_t(key, PyExc_IndexError);
  467. if (index < 0) {
  468. return 0;
  469. }
  470. mkqlKey = NUdf::TUnboxedValuePod(ui64(index));
  471. }
  472. return NUdf::TBoxedValueAccessor::Contains(*dict->Value.Get(), mkqlKey) ? 1 : 0;
  473. } PY_CATCH(-1)
  474. }
  475. PyObject* TPyLazyDict::Get(PyObject* self, PyObject* args)
  476. {
  477. PY_TRY {
  478. PyObject* key = nullptr;
  479. PyObject* failobj = Py_None;
  480. if (!PyArg_UnpackTuple(args, "get", 1, 2, &key, &failobj))
  481. return nullptr;
  482. TPyLazyDict* dict = Cast(self);
  483. if (dict->KeyType) {
  484. const auto mkqlKey = FromPyObject(dict->CastCtx, dict->KeyType, key);
  485. if (auto value = NUdf::TBoxedValueAccessor::Lookup(*dict->Value.Get(), mkqlKey)) {
  486. return ToPyObject(dict->CastCtx, dict->PayloadType, value.Release().GetOptionalValue()).Release();
  487. }
  488. } else {
  489. if (!PyIndex_Check(key)) {
  490. const TPyObjectPtr type = PyObject_Type(key);
  491. const TPyObjectPtr repr = PyObject_Repr(type.Get());
  492. const TPyObjectPtr error = PyUnicode_FromFormat("Unsupported index object type: %R", repr.Get());
  493. PyErr_SetObject(PyExc_TypeError, error.Get());
  494. return nullptr;
  495. }
  496. const Py_ssize_t index = PyNumber_AsSsize_t(key, PyExc_IndexError);
  497. if (index < 0) {
  498. return nullptr;
  499. }
  500. if (auto value = NUdf::TBoxedValueAccessor::Lookup(*dict->Value.Get(), NUdf::TUnboxedValuePod(ui64(index)))) {
  501. return ToPyObject(dict->CastCtx, dict->PayloadType, value.Release().GetOptionalValue()).Release();
  502. }
  503. }
  504. Py_INCREF(failobj);
  505. return failobj;
  506. } PY_CATCH(nullptr)
  507. }
  508. PyObject* TPyLazyDict::Keys(PyObject* self, PyObject* /* args */)
  509. {
  510. PY_TRY {
  511. const auto dict = Cast(self);
  512. return ToPyIterator(dict->CastCtx, dict->KeyType,
  513. NUdf::TBoxedValueAccessor::GetKeysIterator(*dict->Value.Get())).Release();
  514. } PY_CATCH(nullptr)
  515. }
  516. PyObject* TPyLazyDict::Items(PyObject* self, PyObject* /* args */)
  517. {
  518. PY_TRY {
  519. const auto dict = Cast(self);
  520. return ToPyIterator(dict->CastCtx, dict->KeyType, dict->PayloadType,
  521. NUdf::TBoxedValueAccessor::GetDictIterator(*dict->Value.Get())).Release();
  522. } PY_CATCH(nullptr)
  523. }
  524. PyObject* TPyLazyDict::Values(PyObject* self, PyObject* /* args */)
  525. {
  526. PY_TRY {
  527. const auto dict = Cast(self);
  528. return ToPyIterator(dict->CastCtx, dict->PayloadType,
  529. NUdf::TBoxedValueAccessor::GetPayloadsIterator(*dict->Value.Get())).Release();
  530. } PY_CATCH(nullptr)
  531. }
  532. PyObject* TPyLazyDict::New(
  533. const TPyCastContext::TPtr& castCtx,
  534. const NUdf::TType* keyType,
  535. const NUdf::TType* payloadType,
  536. NUdf::IBoxedValuePtr&& value)
  537. {
  538. TPyLazyDict* dict = new TPyLazyDict;
  539. PyObject_INIT(dict, &PyLazyDictType);
  540. dict->CastCtx = castCtx;
  541. dict->KeyType = keyType;
  542. dict->PayloadType = payloadType;
  543. dict->Value.Set(castCtx->PyCtx, value);
  544. return reinterpret_cast<PyObject*>(dict);
  545. }
  546. //////////////////////////////////////////////////////////////////////////////
  547. // TPyLazySet implementation
  548. //////////////////////////////////////////////////////////////////////////////
  549. int TPyLazySet::Bool(PyObject* self)
  550. {
  551. PY_TRY {
  552. return NUdf::TBoxedValueAccessor::HasDictItems(*Cast(self)->Value.Get()) ? 1 : 0;
  553. } PY_CATCH(-1)
  554. }
  555. PyObject* TPyLazySet::Repr(PyObject*)
  556. {
  557. return PyRepr("<yql.TSet>").Release();
  558. }
  559. Py_ssize_t TPyLazySet::Len(PyObject* self)
  560. {
  561. PY_TRY {
  562. return static_cast<Py_ssize_t>(NUdf::TBoxedValueAccessor::GetDictLength(*Cast(self)->Value.Get()));
  563. } PY_CATCH(-1)
  564. }
  565. // -1 error
  566. // 0 not found
  567. // 1 found
  568. int TPyLazySet::Contains(PyObject* self, PyObject* key)
  569. {
  570. PY_TRY {
  571. const auto set = Cast(self);
  572. const auto mkqlKey = FromPyObject(set->CastCtx, set->ItemType, key);
  573. return NUdf::TBoxedValueAccessor::Contains(*set->Value.Get(), mkqlKey) ? 1 : 0;
  574. } PY_CATCH(-1)
  575. }
  576. PyObject* TPyLazySet::Iter(PyObject* self)
  577. {
  578. PY_TRY {
  579. const auto set = Cast(self);
  580. return ToPyIterator(set->CastCtx, set->ItemType,
  581. NUdf::TBoxedValueAccessor::GetKeysIterator(*set->Value.Get())).Release();
  582. } PY_CATCH(nullptr)
  583. }
  584. PyObject* TPyLazySet::New(
  585. const TPyCastContext::TPtr& castCtx,
  586. const NUdf::TType* itemType,
  587. NUdf::IBoxedValuePtr&& value)
  588. {
  589. TPyLazySet* dict = new TPyLazySet;
  590. PyObject_INIT(dict, &PyLazySetType);
  591. dict->CastCtx = castCtx;
  592. dict->ItemType = itemType;
  593. dict->Value.Set(castCtx->PyCtx, value);
  594. return reinterpret_cast<PyObject*>(dict);
  595. }
  596. //////////////////////////////////////////////////////////////////////////////
  597. TPyObjectPtr ToPyLazyDict(
  598. const TPyCastContext::TPtr& castCtx,
  599. const NUdf::TType* keyType,
  600. const NUdf::TType* payloadType,
  601. const NUdf::TUnboxedValuePod& value)
  602. {
  603. return TPyLazyDict::New(castCtx, keyType, payloadType, value.AsBoxed());
  604. }
  605. TPyObjectPtr ToPyLazySet(
  606. const TPyCastContext::TPtr& castCtx,
  607. const NUdf::TType* itemType,
  608. const NUdf::TUnboxedValuePod& value)
  609. {
  610. return TPyLazySet::New(castCtx, itemType, value.AsBoxed());
  611. }
  612. } // namspace NPython