generalized str::operator std::string() to accept 'bytes'(3.x)/'string'(2.7)

wjakob · wjakob · commit 5612a0c1c226 · 2016-05-01T00:32:18.000+02:00
diff --git a/include/pybind11/common.h b/include/pybind11/common.h
@@ -81,6 +81,7 @@
 #define PYBIND11_BYTES_FROM_STRING_AND_SIZE PyBytes_FromStringAndSize
 #define PYBIND11_BYTES_AS_STRING_AND_SIZE PyBytes_AsStringAndSize
 #define PYBIND11_BYTES_AS_STRING PyBytes_AsString
+#define PYBIND11_BYTES_CHECK PyBytes_Check
 #define PYBIND11_LONG_CHECK(o) PyLong_Check(o)
 #define PYBIND11_LONG_AS_LONGLONG(o) PyLong_AsLongLong(o)
 #define PYBIND11_LONG_AS_UNSIGNED_LONGLONG(o) PyLong_AsUnsignedLongLong(o)
@@ -98,6 +99,7 @@
 #define PYBIND11_BYTES_FROM_STRING_AND_SIZE PyString_FromStringAndSize
 #define PYBIND11_BYTES_AS_STRING_AND_SIZE PyString_AsStringAndSize
 #define PYBIND11_BYTES_AS_STRING PyString_AsString
+#define PYBIND11_BYTES_CHECK PyString_Check
 #define PYBIND11_LONG_CHECK(o) (PyInt_Check(o) || PyLong_Check(o))
 #define PYBIND11_LONG_AS_LONGLONG(o) (PyInt_Check(o) ? (long long) PyLong_AsLong(o) : PyLong_AsLongLong(o))
 #define PYBIND11_LONG_AS_UNSIGNED_LONGLONG(o) (PyInt_Check(o) ? (unsigned long long) PyLong_AsUnsignedLong(o) : PyLong_AsUnsignedLongLong(o))
diff --git a/include/pybind11/pytypes.h b/include/pybind11/pytypes.h
@@ -163,7 +163,7 @@ struct list_accessor {
         return object(result, true);
     }
 
-    template <typename T> inline T cast() const { return operator object().cast<T>(); }
+    template <typename T> T cast() const { return operator object().cast<T>(); }
 private:
     handle list;
     size_t index;
@@ -188,7 +188,7 @@ struct tuple_accessor {
         return object(result, true);
     }
 
-    template <typename T> inline T cast() const { return operator object().cast<T>(); }
+    template <typename T> T cast() const { return operator object().cast<T>(); }
 private:
     handle tuple;
     size_t index;
@@ -225,6 +225,8 @@ inline bool PyIterable_Check(PyObject *obj) {
     
 inline bool PyNone_Check(PyObject *o) { return o == Py_None; }
 
+inline bool PyUnicode_Check_Permissive(PyObject *o) { return PyUnicode_Check(o) || PYBIND11_BYTES_CHECK(o); }
+
 NAMESPACE_END(detail)
 
 #define PYBIND11_OBJECT_CVT(Name, Parent, CheckFun, CvtStmt) \
@@ -316,21 +318,26 @@ inline iterator handle::end() const { return iterator(nullptr, false); }
 
 class str : public object {
 public:
-    PYBIND11_OBJECT_DEFAULT(str, object, PyUnicode_Check)
+    PYBIND11_OBJECT_DEFAULT(str, object, detail::PyUnicode_Check_Permissive)
+
     str(const std::string &s)
         : object(PyUnicode_FromStringAndSize(s.c_str(), s.length()), false) {
         if (!m_ptr) pybind11_fail("Could not allocate string object!");
     }
 
     operator std::string() const {
-#if PY_MAJOR_VERSION >= 3 && PY_MINOR_VERSION >= 3
-        return PyUnicode_AsUTF8(m_ptr);
-#else
-        object temp(PyUnicode_AsUTF8String(m_ptr), false);
-        if (temp.ptr() == nullptr)
-            pybind11_fail("Unable to extract string contents!");
-        return PYBIND11_BYTES_AS_STRING(temp.ptr());
-#endif
+        object temp = *this;
+        if (PyUnicode_Check(m_ptr)) {
+            temp = object(PyUnicode_AsUTF8String(m_ptr), false);
+            if (!temp)
+                pybind11_fail("Unable to extract string contents! (encoding issue)");
+        }
+        char *buffer;
+        ssize_t length;
+        int err = PYBIND11_BYTES_AS_STRING_AND_SIZE(temp.ptr(), &buffer, &length);
+        if (err == -1)
+            pybind11_fail("Unable to extract string contents! (invalid type)");
+        return std::string(buffer, length);
     }
 };
 
