From 31a4c55e768bc6358bda4529672808f71abf686b Mon Sep 17 00:00:00 2001
From: Teon Banek <teon.banek@memgraph.io>
Date: Wed, 4 Mar 2020 12:48:02 +0100
Subject: [PATCH] Add HasAttr, GetAttr, and SetAttr to py::Object

Reviewers: ipaljak

Reviewed By: ipaljak

Subscribers: pullbot

Differential Revision: https://phabricator.memgraph.io/D2706
---
 src/py/py.hpp                     | 55 +++++++++++++++++++++++++++++--
 src/query/procedure/py_module.cpp |  8 ++---
 2 files changed, 56 insertions(+), 7 deletions(-)

diff --git a/src/py/py.hpp b/src/py/py.hpp
index d88631575..73b20e6c3 100644
--- a/src/py/py.hpp
+++ b/src/py/py.hpp
@@ -38,7 +38,7 @@ class EnsureGIL final {
 };
 
 /// Owns a `PyObject *` and supports a more C++ idiomatic API to objects.
-class Object final {
+class [[nodiscard]] Object final {
   PyObject *ptr_{nullptr};
 
  public:
@@ -85,6 +85,56 @@ class Object final {
   /// @sa FetchError
   Object Str() const { return Object(PyObject_Str(ptr_)); }
 
+  /// Equivalent to `hasattr(this, attr_name)` in Python.
+  ///
+  /// This function always succeeds, meaning that exceptions that occur while
+  /// calling __getattr__ and __getattribute__ will get suppressed. To get error
+  /// reporting, use GetAttr instead.
+  bool HasAttr(const char *attr_name) const {
+    return PyObject_HasAttrString(ptr_, attr_name);
+  }
+
+  /// Equivalent to `hasattr(this, attr_name)` in Python.
+  ///
+  /// This function always succeeds, meaning that exceptions that occur while
+  /// calling __getattr__ and __getattribute__ will get suppressed. To get error
+  /// reporting, use GetAttr instead.
+  bool HasAttr(PyObject *attr_name) const {
+    return PyObject_HasAttr(ptr_, attr_name);
+  }
+
+  /// Equivalent to `this.attr_name` in Python.
+  ///
+  /// Returned Object is nullptr if an error occurred.
+  /// @sa FetchError
+  Object GetAttr(const char *attr_name) const {
+    return Object(PyObject_GetAttrString(ptr_, attr_name));
+  }
+
+  /// Equivalent to `this.attr_name` in Python.
+  ///
+  /// Returned Object is nullptr if an error occurred.
+  /// @sa FetchError
+  Object GetAttr(PyObject *attr_name) const {
+    return Object(PyObject_GetAttr(ptr_, attr_name));
+  }
+
+  /// Equivalent to `this.attr_name = v` in Python.
+  ///
+  /// False is returned if an error occurred.
+  /// @sa FetchError
+  [[nodiscard]] bool SetAttr(const char *attr_name, PyObject *v) {
+    return PyObject_SetAttrString(ptr_, attr_name, v) == 0;
+  }
+
+  /// Equivalent to `this.attr_name = v` in Python.
+  ///
+  /// False is returned if an error occurred.
+  /// @sa FetchError
+  [[nodiscard]] bool SetAttr(PyObject *attr_name, PyObject *v) {
+    return PyObject_SetAttr(ptr_, attr_name, v) == 0;
+  }
+
   /// Equivalent to `callable()` in Python.
   ///
   /// Returned Object is nullptr if an error occurred.
@@ -148,8 +198,7 @@ inline std::ostream &operator<<(std::ostream &os,
   if (!exc_info.type) return os;
   Object traceback_mod(PyImport_ImportModule("traceback"));
   CHECK(traceback_mod);
-  Object format_exception_fn(
-      PyObject_GetAttrString(traceback_mod, "format_exception"));
+  Object format_exception_fn(traceback_mod.GetAttr("format_exception"));
   CHECK(format_exception_fn);
   auto list = format_exception_fn.Call(
       exc_info.type, exc_info.value ? exc_info.value : Py_None,
diff --git a/src/query/procedure/py_module.cpp b/src/query/procedure/py_module.cpp
index 51697a7b4..509a966fb 100644
--- a/src/query/procedure/py_module.cpp
+++ b/src/query/procedure/py_module.cpp
@@ -411,7 +411,7 @@ PyObject *PyQueryModuleAddReadProcedure(PyQueryModule *self, PyObject *cb) {
   }
   Py_INCREF(cb);
   py::Object py_cb(cb);
-  py::Object py_name(PyObject_GetAttrString(py_cb, "__name__"));
+  py::Object py_name(py_cb.GetAttr("__name__"));
   const auto *name = PyUnicode_AsUTF8(py_name);
   // TODO: Validate name
   auto *memory = self->module->procedures.get_allocator().GetMemoryResource();
@@ -1055,7 +1055,7 @@ template <class TFun>
 auto WithMgpModule(mgp_module *module_def, const TFun &fun) {
   py::Object py_mgp(PyImport_ImportModule("_mgp"));
   CHECK(py_mgp) << "Expected builtin '_mgp' to be available for import";
-  py::Object py_mgp_module(PyObject_GetAttrString(py_mgp, "_MODULE"));
+  py::Object py_mgp_module(py_mgp.GetAttr("_MODULE"));
   CHECK(py_mgp_module) << "Expected '_mgp' to have attribute '_MODULE'";
   // NOTE: This check is not thread safe, but this should only go through
   // ModuleRegistry::LoadModuleLibrary which ought to serialize loading.
@@ -1065,9 +1065,9 @@ auto WithMgpModule(mgp_module *module_def, const TFun &fun) {
          "modules?";
   auto *py_query_module = MakePyQueryModule(module_def);
   CHECK(py_query_module);
-  CHECK(0 <= PyObject_SetAttrString(py_mgp, "_MODULE", py_query_module));
+  CHECK(py_mgp.SetAttr("_MODULE", py_query_module));
   auto ret = fun();
-  CHECK(0 <= PyObject_SetAttrString(py_mgp, "_MODULE", Py_None));
+  CHECK(py_mgp.SetAttr("_MODULE", Py_None));
   return ret;
 }