店播爬取Python脚本

map_container.cc 37KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060
  1. // Protocol Buffers - Google's data interchange format
  2. // Copyright 2008 Google Inc. All rights reserved.
  3. // https://developers.google.com/protocol-buffers/
  4. //
  5. // Redistribution and use in source and binary forms, with or without
  6. // modification, are permitted provided that the following conditions are
  7. // met:
  8. //
  9. // * Redistributions of source code must retain the above copyright
  10. // notice, this list of conditions and the following disclaimer.
  11. // * Redistributions in binary form must reproduce the above
  12. // copyright notice, this list of conditions and the following disclaimer
  13. // in the documentation and/or other materials provided with the
  14. // distribution.
  15. // * Neither the name of Google Inc. nor the names of its
  16. // contributors may be used to endorse or promote products derived from
  17. // this software without specific prior written permission.
  18. //
  19. // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  20. // "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  21. // LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  22. // A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  23. // OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  24. // SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  25. // LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  26. // DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  27. // THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  28. // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  29. // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  30. // Author: haberman@google.com (Josh Haberman)
  31. #include <google/protobuf/pyext/map_container.h>
  32. #include <cstdint>
  33. #include <memory>
  34. #include <google/protobuf/stubs/logging.h>
  35. #include <google/protobuf/stubs/common.h>
  36. #include <google/protobuf/map.h>
  37. #include <google/protobuf/map_field.h>
  38. #include <google/protobuf/message.h>
  39. #include <google/protobuf/pyext/message.h>
  40. #include <google/protobuf/pyext/message_factory.h>
  41. #include <google/protobuf/pyext/repeated_composite_container.h>
  42. #include <google/protobuf/pyext/scoped_pyobject_ptr.h>
  43. #include <google/protobuf/stubs/map_util.h>
  44. #if PY_MAJOR_VERSION >= 3
  45. #define PyInt_FromLong PyLong_FromLong
  46. #define PyInt_FromSize_t PyLong_FromSize_t
  47. #endif
  48. namespace google {
  49. namespace protobuf {
  50. namespace python {
  51. // Functions that need access to map reflection functionality.
  52. // They need to be contained in this class because it is friended.
  53. class MapReflectionFriend {
  54. public:
  55. // Methods that are in common between the map types.
  56. static PyObject* Contains(PyObject* _self, PyObject* key);
  57. static Py_ssize_t Length(PyObject* _self);
  58. static PyObject* GetIterator(PyObject *_self);
  59. static PyObject* IterNext(PyObject* _self);
  60. static PyObject* MergeFrom(PyObject* _self, PyObject* arg);
  61. // Methods that differ between the map types.
  62. static PyObject* ScalarMapGetItem(PyObject* _self, PyObject* key);
  63. static PyObject* MessageMapGetItem(PyObject* _self, PyObject* key);
  64. static int ScalarMapSetItem(PyObject* _self, PyObject* key, PyObject* v);
  65. static int MessageMapSetItem(PyObject* _self, PyObject* key, PyObject* v);
  66. static PyObject* ScalarMapToStr(PyObject* _self);
  67. static PyObject* MessageMapToStr(PyObject* _self);
  68. };
  69. struct MapIterator {
  70. PyObject_HEAD;
  71. std::unique_ptr<::google::protobuf::MapIterator> iter;
  72. // A pointer back to the container, so we can notice changes to the version.
  73. // We own a ref on this.
  74. MapContainer* container;
  75. // We need to keep a ref on the parent Message too, because
  76. // MapIterator::~MapIterator() accesses it. Normally this would be ok because
  77. // the ref on container (above) would guarantee outlive semantics. However in
  78. // the case of ClearField(), the MapContainer points to a different message,
  79. // a copy of the original. But our iterator still points to the original,
  80. // which could now get deleted before us.
  81. //
  82. // To prevent this, we ensure that the Message will always stay alive as long
  83. // as this iterator does. This is solely for the benefit of the MapIterator
  84. // destructor -- we should never actually access the iterator in this state
  85. // except to delete it.
  86. CMessage* parent;
  87. // The version of the map when we took the iterator to it.
  88. //
  89. // We store this so that if the map is modified during iteration we can throw
  90. // an error.
  91. uint64_t version;
  92. };
  93. Message* MapContainer::GetMutableMessage() {
  94. cmessage::AssureWritable(parent);
  95. return parent->message;
  96. }
  97. // Consumes a reference on the Python string object.
  98. static bool PyStringToSTL(PyObject* py_string, std::string* stl_string) {
  99. char *value;
  100. Py_ssize_t value_len;
  101. if (!py_string) {
  102. return false;
  103. }
  104. if (PyBytes_AsStringAndSize(py_string, &value, &value_len) < 0) {
  105. Py_DECREF(py_string);
  106. return false;
  107. } else {
  108. stl_string->assign(value, value_len);
  109. Py_DECREF(py_string);
  110. return true;
  111. }
  112. }
  113. static bool PythonToMapKey(MapContainer* self, PyObject* obj, MapKey* key) {
  114. const FieldDescriptor* field_descriptor =
  115. self->parent_field_descriptor->message_type()->map_key();
  116. switch (field_descriptor->cpp_type()) {
  117. case FieldDescriptor::CPPTYPE_INT32: {
  118. GOOGLE_CHECK_GET_INT32(obj, value, false);
  119. key->SetInt32Value(value);
  120. break;
  121. }
  122. case FieldDescriptor::CPPTYPE_INT64: {
  123. GOOGLE_CHECK_GET_INT64(obj, value, false);
  124. key->SetInt64Value(value);
  125. break;
  126. }
  127. case FieldDescriptor::CPPTYPE_UINT32: {
  128. GOOGLE_CHECK_GET_UINT32(obj, value, false);
  129. key->SetUInt32Value(value);
  130. break;
  131. }
  132. case FieldDescriptor::CPPTYPE_UINT64: {
  133. GOOGLE_CHECK_GET_UINT64(obj, value, false);
  134. key->SetUInt64Value(value);
  135. break;
  136. }
  137. case FieldDescriptor::CPPTYPE_BOOL: {
  138. GOOGLE_CHECK_GET_BOOL(obj, value, false);
  139. key->SetBoolValue(value);
  140. break;
  141. }
  142. case FieldDescriptor::CPPTYPE_STRING: {
  143. std::string str;
  144. if (!PyStringToSTL(CheckString(obj, field_descriptor), &str)) {
  145. return false;
  146. }
  147. key->SetStringValue(str);
  148. break;
  149. }
  150. default:
  151. PyErr_Format(
  152. PyExc_SystemError, "Type %d cannot be a map key",
  153. field_descriptor->cpp_type());
  154. return false;
  155. }
  156. return true;
  157. }
  158. static PyObject* MapKeyToPython(MapContainer* self, const MapKey& key) {
  159. const FieldDescriptor* field_descriptor =
  160. self->parent_field_descriptor->message_type()->map_key();
  161. switch (field_descriptor->cpp_type()) {
  162. case FieldDescriptor::CPPTYPE_INT32:
  163. return PyInt_FromLong(key.GetInt32Value());
  164. case FieldDescriptor::CPPTYPE_INT64:
  165. return PyLong_FromLongLong(key.GetInt64Value());
  166. case FieldDescriptor::CPPTYPE_UINT32:
  167. return PyInt_FromSize_t(key.GetUInt32Value());
  168. case FieldDescriptor::CPPTYPE_UINT64:
  169. return PyLong_FromUnsignedLongLong(key.GetUInt64Value());
  170. case FieldDescriptor::CPPTYPE_BOOL:
  171. return PyBool_FromLong(key.GetBoolValue());
  172. case FieldDescriptor::CPPTYPE_STRING:
  173. return ToStringObject(field_descriptor, key.GetStringValue());
  174. default:
  175. PyErr_Format(
  176. PyExc_SystemError, "Couldn't convert type %d to value",
  177. field_descriptor->cpp_type());
  178. return NULL;
  179. }
  180. }
  181. // This is only used for ScalarMap, so we don't need to handle the
  182. // CPPTYPE_MESSAGE case.
  183. PyObject* MapValueRefToPython(MapContainer* self, const MapValueRef& value) {
  184. const FieldDescriptor* field_descriptor =
  185. self->parent_field_descriptor->message_type()->map_value();
  186. switch (field_descriptor->cpp_type()) {
  187. case FieldDescriptor::CPPTYPE_INT32:
  188. return PyInt_FromLong(value.GetInt32Value());
  189. case FieldDescriptor::CPPTYPE_INT64:
  190. return PyLong_FromLongLong(value.GetInt64Value());
  191. case FieldDescriptor::CPPTYPE_UINT32:
  192. return PyInt_FromSize_t(value.GetUInt32Value());
  193. case FieldDescriptor::CPPTYPE_UINT64:
  194. return PyLong_FromUnsignedLongLong(value.GetUInt64Value());
  195. case FieldDescriptor::CPPTYPE_FLOAT:
  196. return PyFloat_FromDouble(value.GetFloatValue());
  197. case FieldDescriptor::CPPTYPE_DOUBLE:
  198. return PyFloat_FromDouble(value.GetDoubleValue());
  199. case FieldDescriptor::CPPTYPE_BOOL:
  200. return PyBool_FromLong(value.GetBoolValue());
  201. case FieldDescriptor::CPPTYPE_STRING:
  202. return ToStringObject(field_descriptor, value.GetStringValue());
  203. case FieldDescriptor::CPPTYPE_ENUM:
  204. return PyInt_FromLong(value.GetEnumValue());
  205. default:
  206. PyErr_Format(
  207. PyExc_SystemError, "Couldn't convert type %d to value",
  208. field_descriptor->cpp_type());
  209. return NULL;
  210. }
  211. }
  212. // This is only used for ScalarMap, so we don't need to handle the
  213. // CPPTYPE_MESSAGE case.
  214. static bool PythonToMapValueRef(MapContainer* self, PyObject* obj,
  215. bool allow_unknown_enum_values,
  216. MapValueRef* value_ref) {
  217. const FieldDescriptor* field_descriptor =
  218. self->parent_field_descriptor->message_type()->map_value();
  219. switch (field_descriptor->cpp_type()) {
  220. case FieldDescriptor::CPPTYPE_INT32: {
  221. GOOGLE_CHECK_GET_INT32(obj, value, false);
  222. value_ref->SetInt32Value(value);
  223. return true;
  224. }
  225. case FieldDescriptor::CPPTYPE_INT64: {
  226. GOOGLE_CHECK_GET_INT64(obj, value, false);
  227. value_ref->SetInt64Value(value);
  228. return true;
  229. }
  230. case FieldDescriptor::CPPTYPE_UINT32: {
  231. GOOGLE_CHECK_GET_UINT32(obj, value, false);
  232. value_ref->SetUInt32Value(value);
  233. return true;
  234. }
  235. case FieldDescriptor::CPPTYPE_UINT64: {
  236. GOOGLE_CHECK_GET_UINT64(obj, value, false);
  237. value_ref->SetUInt64Value(value);
  238. return true;
  239. }
  240. case FieldDescriptor::CPPTYPE_FLOAT: {
  241. GOOGLE_CHECK_GET_FLOAT(obj, value, false);
  242. value_ref->SetFloatValue(value);
  243. return true;
  244. }
  245. case FieldDescriptor::CPPTYPE_DOUBLE: {
  246. GOOGLE_CHECK_GET_DOUBLE(obj, value, false);
  247. value_ref->SetDoubleValue(value);
  248. return true;
  249. }
  250. case FieldDescriptor::CPPTYPE_BOOL: {
  251. GOOGLE_CHECK_GET_BOOL(obj, value, false);
  252. value_ref->SetBoolValue(value);
  253. return true;;
  254. }
  255. case FieldDescriptor::CPPTYPE_STRING: {
  256. std::string str;
  257. if (!PyStringToSTL(CheckString(obj, field_descriptor), &str)) {
  258. return false;
  259. }
  260. value_ref->SetStringValue(str);
  261. return true;
  262. }
  263. case FieldDescriptor::CPPTYPE_ENUM: {
  264. GOOGLE_CHECK_GET_INT32(obj, value, false);
  265. if (allow_unknown_enum_values) {
  266. value_ref->SetEnumValue(value);
  267. return true;
  268. } else {
  269. const EnumDescriptor* enum_descriptor = field_descriptor->enum_type();
  270. const EnumValueDescriptor* enum_value =
  271. enum_descriptor->FindValueByNumber(value);
  272. if (enum_value != NULL) {
  273. value_ref->SetEnumValue(value);
  274. return true;
  275. } else {
  276. PyErr_Format(PyExc_ValueError, "Unknown enum value: %d", value);
  277. return false;
  278. }
  279. }
  280. break;
  281. }
  282. default:
  283. PyErr_Format(
  284. PyExc_SystemError, "Setting value to a field of unknown type %d",
  285. field_descriptor->cpp_type());
  286. return false;
  287. }
  288. }
  289. // Map methods common to ScalarMap and MessageMap //////////////////////////////
  290. static MapContainer* GetMap(PyObject* obj) {
  291. return reinterpret_cast<MapContainer*>(obj);
  292. }
  293. Py_ssize_t MapReflectionFriend::Length(PyObject* _self) {
  294. MapContainer* self = GetMap(_self);
  295. const google::protobuf::Message* message = self->parent->message;
  296. return message->GetReflection()->MapSize(*message,
  297. self->parent_field_descriptor);
  298. }
  299. PyObject* Clear(PyObject* _self) {
  300. MapContainer* self = GetMap(_self);
  301. Message* message = self->GetMutableMessage();
  302. const Reflection* reflection = message->GetReflection();
  303. reflection->ClearField(message, self->parent_field_descriptor);
  304. Py_RETURN_NONE;
  305. }
  306. PyObject* GetEntryClass(PyObject* _self) {
  307. MapContainer* self = GetMap(_self);
  308. CMessageClass* message_class = message_factory::GetMessageClass(
  309. cmessage::GetFactoryForMessage(self->parent),
  310. self->parent_field_descriptor->message_type());
  311. Py_XINCREF(message_class);
  312. return reinterpret_cast<PyObject*>(message_class);
  313. }
  314. PyObject* MapReflectionFriend::MergeFrom(PyObject* _self, PyObject* arg) {
  315. MapContainer* self = GetMap(_self);
  316. if (!PyObject_TypeCheck(arg, ScalarMapContainer_Type) &&
  317. !PyObject_TypeCheck(arg, MessageMapContainer_Type)) {
  318. PyErr_SetString(PyExc_AttributeError, "Not a map field");
  319. return nullptr;
  320. }
  321. MapContainer* other_map = GetMap(arg);
  322. Message* message = self->GetMutableMessage();
  323. const Message* other_message = other_map->parent->message;
  324. const Reflection* reflection = message->GetReflection();
  325. const Reflection* other_reflection = other_message->GetReflection();
  326. internal::MapFieldBase* field = reflection->MutableMapData(
  327. message, self->parent_field_descriptor);
  328. const internal::MapFieldBase* other_field = other_reflection->GetMapData(
  329. *other_message, other_map->parent_field_descriptor);
  330. field->MergeFrom(*other_field);
  331. self->version++;
  332. Py_RETURN_NONE;
  333. }
  334. PyObject* MapReflectionFriend::Contains(PyObject* _self, PyObject* key) {
  335. MapContainer* self = GetMap(_self);
  336. const Message* message = self->parent->message;
  337. const Reflection* reflection = message->GetReflection();
  338. MapKey map_key;
  339. if (!PythonToMapKey(self, key, &map_key)) {
  340. return NULL;
  341. }
  342. if (reflection->ContainsMapKey(*message, self->parent_field_descriptor,
  343. map_key)) {
  344. Py_RETURN_TRUE;
  345. } else {
  346. Py_RETURN_FALSE;
  347. }
  348. }
  349. // ScalarMap ///////////////////////////////////////////////////////////////////
  350. MapContainer* NewScalarMapContainer(
  351. CMessage* parent, const google::protobuf::FieldDescriptor* parent_field_descriptor) {
  352. if (!CheckFieldBelongsToMessage(parent_field_descriptor, parent->message)) {
  353. return NULL;
  354. }
  355. PyObject* obj(PyType_GenericAlloc(ScalarMapContainer_Type, 0));
  356. if (obj == NULL) {
  357. PyErr_Format(PyExc_RuntimeError,
  358. "Could not allocate new container.");
  359. return NULL;
  360. }
  361. MapContainer* self = GetMap(obj);
  362. Py_INCREF(parent);
  363. self->parent = parent;
  364. self->parent_field_descriptor = parent_field_descriptor;
  365. self->version = 0;
  366. return self;
  367. }
  368. PyObject* MapReflectionFriend::ScalarMapGetItem(PyObject* _self,
  369. PyObject* key) {
  370. MapContainer* self = GetMap(_self);
  371. Message* message = self->GetMutableMessage();
  372. const Reflection* reflection = message->GetReflection();
  373. MapKey map_key;
  374. MapValueRef value;
  375. if (!PythonToMapKey(self, key, &map_key)) {
  376. return NULL;
  377. }
  378. if (reflection->InsertOrLookupMapValue(message, self->parent_field_descriptor,
  379. map_key, &value)) {
  380. self->version++;
  381. }
  382. return MapValueRefToPython(self, value);
  383. }
  384. int MapReflectionFriend::ScalarMapSetItem(PyObject* _self, PyObject* key,
  385. PyObject* v) {
  386. MapContainer* self = GetMap(_self);
  387. Message* message = self->GetMutableMessage();
  388. const Reflection* reflection = message->GetReflection();
  389. MapKey map_key;
  390. MapValueRef value;
  391. if (!PythonToMapKey(self, key, &map_key)) {
  392. return -1;
  393. }
  394. self->version++;
  395. if (v) {
  396. // Set item to v.
  397. reflection->InsertOrLookupMapValue(message, self->parent_field_descriptor,
  398. map_key, &value);
  399. if (!PythonToMapValueRef(self, v, reflection->SupportsUnknownEnumValues(),
  400. &value)) {
  401. return -1;
  402. }
  403. return 0;
  404. } else {
  405. // Delete key from map.
  406. if (reflection->DeleteMapValue(message, self->parent_field_descriptor,
  407. map_key)) {
  408. return 0;
  409. } else {
  410. PyErr_Format(PyExc_KeyError, "Key not present in map");
  411. return -1;
  412. }
  413. }
  414. }
  415. static PyObject* ScalarMapGet(PyObject* self, PyObject* args,
  416. PyObject* kwargs) {
  417. static const char* kwlist[] = {"key", "default", nullptr};
  418. PyObject* key;
  419. PyObject* default_value = NULL;
  420. if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|O",
  421. const_cast<char**>(kwlist), &key,
  422. &default_value)) {
  423. return NULL;
  424. }
  425. ScopedPyObjectPtr is_present(MapReflectionFriend::Contains(self, key));
  426. if (is_present.get() == NULL) {
  427. return NULL;
  428. }
  429. if (PyObject_IsTrue(is_present.get())) {
  430. return MapReflectionFriend::ScalarMapGetItem(self, key);
  431. } else {
  432. if (default_value != NULL) {
  433. Py_INCREF(default_value);
  434. return default_value;
  435. } else {
  436. Py_RETURN_NONE;
  437. }
  438. }
  439. }
  440. PyObject* MapReflectionFriend::ScalarMapToStr(PyObject* _self) {
  441. ScopedPyObjectPtr dict(PyDict_New());
  442. if (dict == NULL) {
  443. return NULL;
  444. }
  445. ScopedPyObjectPtr key;
  446. ScopedPyObjectPtr value;
  447. MapContainer* self = GetMap(_self);
  448. Message* message = self->GetMutableMessage();
  449. const Reflection* reflection = message->GetReflection();
  450. for (google::protobuf::MapIterator it = reflection->MapBegin(
  451. message, self->parent_field_descriptor);
  452. it != reflection->MapEnd(message, self->parent_field_descriptor);
  453. ++it) {
  454. key.reset(MapKeyToPython(self, it.GetKey()));
  455. if (key == NULL) {
  456. return NULL;
  457. }
  458. value.reset(MapValueRefToPython(self, it.GetValueRef()));
  459. if (value == NULL) {
  460. return NULL;
  461. }
  462. if (PyDict_SetItem(dict.get(), key.get(), value.get()) < 0) {
  463. return NULL;
  464. }
  465. }
  466. return PyObject_Repr(dict.get());
  467. }
  468. static void ScalarMapDealloc(PyObject* _self) {
  469. MapContainer* self = GetMap(_self);
  470. self->RemoveFromParentCache();
  471. PyTypeObject *type = Py_TYPE(_self);
  472. type->tp_free(_self);
  473. if (type->tp_flags & Py_TPFLAGS_HEAPTYPE) {
  474. // With Python3, the Map class is not static, and must be managed.
  475. Py_DECREF(type);
  476. }
  477. }
  478. static PyMethodDef ScalarMapMethods[] = {
  479. {"__contains__", MapReflectionFriend::Contains, METH_O,
  480. "Tests whether a key is a member of the map."},
  481. {"clear", (PyCFunction)Clear, METH_NOARGS,
  482. "Removes all elements from the map."},
  483. {"get", (PyCFunction)ScalarMapGet, METH_VARARGS | METH_KEYWORDS,
  484. "Gets the value for the given key if present, or otherwise a default"},
  485. {"GetEntryClass", (PyCFunction)GetEntryClass, METH_NOARGS,
  486. "Return the class used to build Entries of (key, value) pairs."},
  487. {"MergeFrom", (PyCFunction)MapReflectionFriend::MergeFrom, METH_O,
  488. "Merges a map into the current map."},
  489. /*
  490. { "__deepcopy__", (PyCFunction)DeepCopy, METH_VARARGS,
  491. "Makes a deep copy of the class." },
  492. { "__reduce__", (PyCFunction)Reduce, METH_NOARGS,
  493. "Outputs picklable representation of the repeated field." },
  494. */
  495. {NULL, NULL},
  496. };
  497. PyTypeObject *ScalarMapContainer_Type;
  498. #if PY_MAJOR_VERSION >= 3
  499. static PyType_Slot ScalarMapContainer_Type_slots[] = {
  500. {Py_tp_dealloc, (void *)ScalarMapDealloc},
  501. {Py_mp_length, (void *)MapReflectionFriend::Length},
  502. {Py_mp_subscript, (void *)MapReflectionFriend::ScalarMapGetItem},
  503. {Py_mp_ass_subscript, (void *)MapReflectionFriend::ScalarMapSetItem},
  504. {Py_tp_methods, (void *)ScalarMapMethods},
  505. {Py_tp_iter, (void *)MapReflectionFriend::GetIterator},
  506. {Py_tp_repr, (void *)MapReflectionFriend::ScalarMapToStr},
  507. {0, 0},
  508. };
  509. PyType_Spec ScalarMapContainer_Type_spec = {
  510. FULL_MODULE_NAME ".ScalarMapContainer",
  511. sizeof(MapContainer),
  512. 0,
  513. Py_TPFLAGS_DEFAULT,
  514. ScalarMapContainer_Type_slots
  515. };
  516. #else
  517. static PyMappingMethods ScalarMapMappingMethods = {
  518. MapReflectionFriend::Length, // mp_length
  519. MapReflectionFriend::ScalarMapGetItem, // mp_subscript
  520. MapReflectionFriend::ScalarMapSetItem, // mp_ass_subscript
  521. };
  522. PyTypeObject _ScalarMapContainer_Type = {
  523. PyVarObject_HEAD_INIT(&PyType_Type, 0)
  524. FULL_MODULE_NAME ".ScalarMapContainer", // tp_name
  525. sizeof(MapContainer), // tp_basicsize
  526. 0, // tp_itemsize
  527. ScalarMapDealloc, // tp_dealloc
  528. 0, // tp_print
  529. 0, // tp_getattr
  530. 0, // tp_setattr
  531. 0, // tp_compare
  532. MapReflectionFriend::ScalarMapToStr, // tp_repr
  533. 0, // tp_as_number
  534. 0, // tp_as_sequence
  535. &ScalarMapMappingMethods, // tp_as_mapping
  536. 0, // tp_hash
  537. 0, // tp_call
  538. 0, // tp_str
  539. 0, // tp_getattro
  540. 0, // tp_setattro
  541. 0, // tp_as_buffer
  542. Py_TPFLAGS_DEFAULT, // tp_flags
  543. "A scalar map container", // tp_doc
  544. 0, // tp_traverse
  545. 0, // tp_clear
  546. 0, // tp_richcompare
  547. 0, // tp_weaklistoffset
  548. MapReflectionFriend::GetIterator, // tp_iter
  549. 0, // tp_iternext
  550. ScalarMapMethods, // tp_methods
  551. 0, // tp_members
  552. 0, // tp_getset
  553. 0, // tp_base
  554. 0, // tp_dict
  555. 0, // tp_descr_get
  556. 0, // tp_descr_set
  557. 0, // tp_dictoffset
  558. 0, // tp_init
  559. };
  560. #endif
  561. // MessageMap //////////////////////////////////////////////////////////////////
  562. static MessageMapContainer* GetMessageMap(PyObject* obj) {
  563. return reinterpret_cast<MessageMapContainer*>(obj);
  564. }
  565. static PyObject* GetCMessage(MessageMapContainer* self, Message* message) {
  566. // Get or create the CMessage object corresponding to this message.
  567. return self->parent
  568. ->BuildSubMessageFromPointer(self->parent_field_descriptor, message,
  569. self->message_class)
  570. ->AsPyObject();
  571. }
  572. MessageMapContainer* NewMessageMapContainer(
  573. CMessage* parent, const google::protobuf::FieldDescriptor* parent_field_descriptor,
  574. CMessageClass* message_class) {
  575. if (!CheckFieldBelongsToMessage(parent_field_descriptor, parent->message)) {
  576. return NULL;
  577. }
  578. PyObject* obj = PyType_GenericAlloc(MessageMapContainer_Type, 0);
  579. if (obj == NULL) {
  580. PyErr_SetString(PyExc_RuntimeError, "Could not allocate new container.");
  581. return NULL;
  582. }
  583. MessageMapContainer* self = GetMessageMap(obj);
  584. Py_INCREF(parent);
  585. self->parent = parent;
  586. self->parent_field_descriptor = parent_field_descriptor;
  587. self->version = 0;
  588. Py_INCREF(message_class);
  589. self->message_class = message_class;
  590. return self;
  591. }
  592. int MapReflectionFriend::MessageMapSetItem(PyObject* _self, PyObject* key,
  593. PyObject* v) {
  594. if (v) {
  595. PyErr_Format(PyExc_ValueError,
  596. "Direct assignment of submessage not allowed");
  597. return -1;
  598. }
  599. // Now we know that this is a delete, not a set.
  600. MessageMapContainer* self = GetMessageMap(_self);
  601. Message* message = self->GetMutableMessage();
  602. const Reflection* reflection = message->GetReflection();
  603. MapKey map_key;
  604. MapValueRef value;
  605. self->version++;
  606. if (!PythonToMapKey(self, key, &map_key)) {
  607. return -1;
  608. }
  609. // Delete key from map.
  610. if (reflection->ContainsMapKey(*message, self->parent_field_descriptor,
  611. map_key)) {
  612. // Delete key from CMessage dict.
  613. MapValueRef value;
  614. reflection->InsertOrLookupMapValue(message, self->parent_field_descriptor,
  615. map_key, &value);
  616. Message* sub_message = value.MutableMessageValue();
  617. // If there is a living weak reference to an item, we "Release" it,
  618. // otherwise we just discard the C++ value.
  619. if (CMessage* released =
  620. self->parent->MaybeReleaseSubMessage(sub_message)) {
  621. Message* msg = released->message;
  622. released->message = msg->New();
  623. msg->GetReflection()->Swap(msg, released->message);
  624. }
  625. // Delete key from map.
  626. reflection->DeleteMapValue(message, self->parent_field_descriptor,
  627. map_key);
  628. return 0;
  629. } else {
  630. PyErr_Format(PyExc_KeyError, "Key not present in map");
  631. return -1;
  632. }
  633. }
  634. PyObject* MapReflectionFriend::MessageMapGetItem(PyObject* _self,
  635. PyObject* key) {
  636. MessageMapContainer* self = GetMessageMap(_self);
  637. Message* message = self->GetMutableMessage();
  638. const Reflection* reflection = message->GetReflection();
  639. MapKey map_key;
  640. MapValueRef value;
  641. if (!PythonToMapKey(self, key, &map_key)) {
  642. return NULL;
  643. }
  644. if (reflection->InsertOrLookupMapValue(message, self->parent_field_descriptor,
  645. map_key, &value)) {
  646. self->version++;
  647. }
  648. return GetCMessage(self, value.MutableMessageValue());
  649. }
  650. PyObject* MapReflectionFriend::MessageMapToStr(PyObject* _self) {
  651. ScopedPyObjectPtr dict(PyDict_New());
  652. if (dict == NULL) {
  653. return NULL;
  654. }
  655. ScopedPyObjectPtr key;
  656. ScopedPyObjectPtr value;
  657. MessageMapContainer* self = GetMessageMap(_self);
  658. Message* message = self->GetMutableMessage();
  659. const Reflection* reflection = message->GetReflection();
  660. for (google::protobuf::MapIterator it = reflection->MapBegin(
  661. message, self->parent_field_descriptor);
  662. it != reflection->MapEnd(message, self->parent_field_descriptor);
  663. ++it) {
  664. key.reset(MapKeyToPython(self, it.GetKey()));
  665. if (key == NULL) {
  666. return NULL;
  667. }
  668. value.reset(GetCMessage(self, it.MutableValueRef()->MutableMessageValue()));
  669. if (value == NULL) {
  670. return NULL;
  671. }
  672. if (PyDict_SetItem(dict.get(), key.get(), value.get()) < 0) {
  673. return NULL;
  674. }
  675. }
  676. return PyObject_Repr(dict.get());
  677. }
  678. PyObject* MessageMapGet(PyObject* self, PyObject* args, PyObject* kwargs) {
  679. static const char* kwlist[] = {"key", "default", nullptr};
  680. PyObject* key;
  681. PyObject* default_value = NULL;
  682. if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|O",
  683. const_cast<char**>(kwlist), &key,
  684. &default_value)) {
  685. return NULL;
  686. }
  687. ScopedPyObjectPtr is_present(MapReflectionFriend::Contains(self, key));
  688. if (is_present.get() == NULL) {
  689. return NULL;
  690. }
  691. if (PyObject_IsTrue(is_present.get())) {
  692. return MapReflectionFriend::MessageMapGetItem(self, key);
  693. } else {
  694. if (default_value != NULL) {
  695. Py_INCREF(default_value);
  696. return default_value;
  697. } else {
  698. Py_RETURN_NONE;
  699. }
  700. }
  701. }
  702. static void MessageMapDealloc(PyObject* _self) {
  703. MessageMapContainer* self = GetMessageMap(_self);
  704. self->RemoveFromParentCache();
  705. Py_DECREF(self->message_class);
  706. PyTypeObject *type = Py_TYPE(_self);
  707. type->tp_free(_self);
  708. if (type->tp_flags & Py_TPFLAGS_HEAPTYPE) {
  709. // With Python3, the Map class is not static, and must be managed.
  710. Py_DECREF(type);
  711. }
  712. }
  713. static PyMethodDef MessageMapMethods[] = {
  714. {"__contains__", (PyCFunction)MapReflectionFriend::Contains, METH_O,
  715. "Tests whether the map contains this element."},
  716. {"clear", (PyCFunction)Clear, METH_NOARGS,
  717. "Removes all elements from the map."},
  718. {"get", (PyCFunction)MessageMapGet, METH_VARARGS | METH_KEYWORDS,
  719. "Gets the value for the given key if present, or otherwise a default"},
  720. {"get_or_create", MapReflectionFriend::MessageMapGetItem, METH_O,
  721. "Alias for getitem, useful to make explicit that the map is mutated."},
  722. {"GetEntryClass", (PyCFunction)GetEntryClass, METH_NOARGS,
  723. "Return the class used to build Entries of (key, value) pairs."},
  724. {"MergeFrom", (PyCFunction)MapReflectionFriend::MergeFrom, METH_O,
  725. "Merges a map into the current map."},
  726. /*
  727. { "__deepcopy__", (PyCFunction)DeepCopy, METH_VARARGS,
  728. "Makes a deep copy of the class." },
  729. { "__reduce__", (PyCFunction)Reduce, METH_NOARGS,
  730. "Outputs picklable representation of the repeated field." },
  731. */
  732. {NULL, NULL},
  733. };
  734. PyTypeObject *MessageMapContainer_Type;
  735. #if PY_MAJOR_VERSION >= 3
  736. static PyType_Slot MessageMapContainer_Type_slots[] = {
  737. {Py_tp_dealloc, (void *)MessageMapDealloc},
  738. {Py_mp_length, (void *)MapReflectionFriend::Length},
  739. {Py_mp_subscript, (void *)MapReflectionFriend::MessageMapGetItem},
  740. {Py_mp_ass_subscript, (void *)MapReflectionFriend::MessageMapSetItem},
  741. {Py_tp_methods, (void *)MessageMapMethods},
  742. {Py_tp_iter, (void *)MapReflectionFriend::GetIterator},
  743. {Py_tp_repr, (void *)MapReflectionFriend::MessageMapToStr},
  744. {0, 0}
  745. };
  746. PyType_Spec MessageMapContainer_Type_spec = {
  747. FULL_MODULE_NAME ".MessageMapContainer",
  748. sizeof(MessageMapContainer),
  749. 0,
  750. Py_TPFLAGS_DEFAULT,
  751. MessageMapContainer_Type_slots
  752. };
  753. #else
  754. static PyMappingMethods MessageMapMappingMethods = {
  755. MapReflectionFriend::Length, // mp_length
  756. MapReflectionFriend::MessageMapGetItem, // mp_subscript
  757. MapReflectionFriend::MessageMapSetItem, // mp_ass_subscript
  758. };
  759. PyTypeObject _MessageMapContainer_Type = {
  760. PyVarObject_HEAD_INIT(&PyType_Type, 0)
  761. FULL_MODULE_NAME ".MessageMapContainer", // tp_name
  762. sizeof(MessageMapContainer), // tp_basicsize
  763. 0, // tp_itemsize
  764. MessageMapDealloc, // tp_dealloc
  765. 0, // tp_print
  766. 0, // tp_getattr
  767. 0, // tp_setattr
  768. 0, // tp_compare
  769. MapReflectionFriend::MessageMapToStr, // tp_repr
  770. 0, // tp_as_number
  771. 0, // tp_as_sequence
  772. &MessageMapMappingMethods, // tp_as_mapping
  773. 0, // tp_hash
  774. 0, // tp_call
  775. 0, // tp_str
  776. 0, // tp_getattro
  777. 0, // tp_setattro
  778. 0, // tp_as_buffer
  779. Py_TPFLAGS_DEFAULT, // tp_flags
  780. "A map container for message", // tp_doc
  781. 0, // tp_traverse
  782. 0, // tp_clear
  783. 0, // tp_richcompare
  784. 0, // tp_weaklistoffset
  785. MapReflectionFriend::GetIterator, // tp_iter
  786. 0, // tp_iternext
  787. MessageMapMethods, // tp_methods
  788. 0, // tp_members
  789. 0, // tp_getset
  790. 0, // tp_base
  791. 0, // tp_dict
  792. 0, // tp_descr_get
  793. 0, // tp_descr_set
  794. 0, // tp_dictoffset
  795. 0, // tp_init
  796. };
  797. #endif
  798. // MapIterator /////////////////////////////////////////////////////////////////
  799. static MapIterator* GetIter(PyObject* obj) {
  800. return reinterpret_cast<MapIterator*>(obj);
  801. }
  802. PyObject* MapReflectionFriend::GetIterator(PyObject *_self) {
  803. MapContainer* self = GetMap(_self);
  804. ScopedPyObjectPtr obj(PyType_GenericAlloc(&MapIterator_Type, 0));
  805. if (obj == NULL) {
  806. return PyErr_Format(PyExc_KeyError, "Could not allocate iterator");
  807. }
  808. MapIterator* iter = GetIter(obj.get());
  809. Py_INCREF(self);
  810. iter->container = self;
  811. iter->version = self->version;
  812. Py_INCREF(self->parent);
  813. iter->parent = self->parent;
  814. if (MapReflectionFriend::Length(_self) > 0) {
  815. Message* message = self->GetMutableMessage();
  816. const Reflection* reflection = message->GetReflection();
  817. iter->iter.reset(new ::google::protobuf::MapIterator(
  818. reflection->MapBegin(message, self->parent_field_descriptor)));
  819. }
  820. return obj.release();
  821. }
  822. PyObject* MapReflectionFriend::IterNext(PyObject* _self) {
  823. MapIterator* self = GetIter(_self);
  824. // This won't catch mutations to the map performed by MergeFrom(); no easy way
  825. // to address that.
  826. if (self->version != self->container->version) {
  827. return PyErr_Format(PyExc_RuntimeError,
  828. "Map modified during iteration.");
  829. }
  830. if (self->parent != self->container->parent) {
  831. return PyErr_Format(PyExc_RuntimeError,
  832. "Map cleared during iteration.");
  833. }
  834. if (self->iter.get() == NULL) {
  835. return NULL;
  836. }
  837. Message* message = self->container->GetMutableMessage();
  838. const Reflection* reflection = message->GetReflection();
  839. if (*self->iter ==
  840. reflection->MapEnd(message, self->container->parent_field_descriptor)) {
  841. return NULL;
  842. }
  843. PyObject* ret = MapKeyToPython(self->container, self->iter->GetKey());
  844. ++(*self->iter);
  845. return ret;
  846. }
  847. static void DeallocMapIterator(PyObject* _self) {
  848. MapIterator* self = GetIter(_self);
  849. self->iter.reset();
  850. Py_CLEAR(self->container);
  851. Py_CLEAR(self->parent);
  852. Py_TYPE(_self)->tp_free(_self);
  853. }
  854. PyTypeObject MapIterator_Type = {
  855. PyVarObject_HEAD_INIT(&PyType_Type, 0)
  856. FULL_MODULE_NAME ".MapIterator", // tp_name
  857. sizeof(MapIterator), // tp_basicsize
  858. 0, // tp_itemsize
  859. DeallocMapIterator, // tp_dealloc
  860. 0, // tp_print
  861. 0, // tp_getattr
  862. 0, // tp_setattr
  863. 0, // tp_compare
  864. 0, // tp_repr
  865. 0, // tp_as_number
  866. 0, // tp_as_sequence
  867. 0, // tp_as_mapping
  868. 0, // tp_hash
  869. 0, // tp_call
  870. 0, // tp_str
  871. 0, // tp_getattro
  872. 0, // tp_setattro
  873. 0, // tp_as_buffer
  874. Py_TPFLAGS_DEFAULT, // tp_flags
  875. "A scalar map iterator", // tp_doc
  876. 0, // tp_traverse
  877. 0, // tp_clear
  878. 0, // tp_richcompare
  879. 0, // tp_weaklistoffset
  880. PyObject_SelfIter, // tp_iter
  881. MapReflectionFriend::IterNext, // tp_iternext
  882. 0, // tp_methods
  883. 0, // tp_members
  884. 0, // tp_getset
  885. 0, // tp_base
  886. 0, // tp_dict
  887. 0, // tp_descr_get
  888. 0, // tp_descr_set
  889. 0, // tp_dictoffset
  890. 0, // tp_init
  891. };
  892. bool InitMapContainers() {
  893. // ScalarMapContainer_Type derives from our MutableMapping type.
  894. ScopedPyObjectPtr containers(PyImport_ImportModule(
  895. "google.protobuf.internal.containers"));
  896. if (containers == NULL) {
  897. return false;
  898. }
  899. ScopedPyObjectPtr mutable_mapping(
  900. PyObject_GetAttrString(containers.get(), "MutableMapping"));
  901. if (mutable_mapping == NULL) {
  902. return false;
  903. }
  904. Py_INCREF(mutable_mapping.get());
  905. #if PY_MAJOR_VERSION >= 3
  906. ScopedPyObjectPtr bases(PyTuple_Pack(1, mutable_mapping.get()));
  907. if (bases == NULL) {
  908. return false;
  909. }
  910. ScalarMapContainer_Type = reinterpret_cast<PyTypeObject*>(
  911. PyType_FromSpecWithBases(&ScalarMapContainer_Type_spec, bases.get()));
  912. #else
  913. _ScalarMapContainer_Type.tp_base =
  914. reinterpret_cast<PyTypeObject*>(mutable_mapping.get());
  915. if (PyType_Ready(&_ScalarMapContainer_Type) < 0) {
  916. return false;
  917. }
  918. ScalarMapContainer_Type = &_ScalarMapContainer_Type;
  919. #endif
  920. if (PyType_Ready(&MapIterator_Type) < 0) {
  921. return false;
  922. }
  923. #if PY_MAJOR_VERSION >= 3
  924. MessageMapContainer_Type = reinterpret_cast<PyTypeObject*>(
  925. PyType_FromSpecWithBases(&MessageMapContainer_Type_spec, bases.get()));
  926. #else
  927. Py_INCREF(mutable_mapping.get());
  928. _MessageMapContainer_Type.tp_base =
  929. reinterpret_cast<PyTypeObject*>(mutable_mapping.get());
  930. if (PyType_Ready(&_MessageMapContainer_Type) < 0) {
  931. return false;
  932. }
  933. MessageMapContainer_Type = &_MessageMapContainer_Type;
  934. #endif
  935. return true;
  936. }
  937. } // namespace python
  938. } // namespace protobuf
  939. } // namespace google