Intermediate changes

commit_hash:0897b988c8ca7235b4f13624177866ac14a46fdd

1 files changed, 51 insertions, 7 deletions

diff --git a/yql/essentials/udfs/common/python/bindings/py_struct.cpp b/yql/essentials/udfs/common/python/bindings/py_struct.cpp
index a4ab99ee32c..bcfd86351fe 100644
--- a/yql/essentials/udfs/common/python/bindings/py_struct.cpp
+++ b/yql/essentials/udfs/common/python/bindings/py_struct.cpp
@@ -72,6 +72,54 @@ TPyObjectPtr CreateNewStrucInstance(const TPyCastContext::TPtr& ctx, const NKiki
     return object;
 }
 
+TPyObjectPtr ConvertStringToPyObject(TStringBuf str)
+{
+#if PY_MAJOR_VERSION >= 3
+    return PyUnicode_FromStringAndSize(str.data(), str.size());
+#else
+    return PyString_FromStringAndSize(str.data(), str.size());
+#endif
+}
+
+// Sized counterpart of PyDict_GetItemString from C API. Returns borrowed reference.
+PyObject* GetItemFromPyDictUsingString(PyObject* v, TStringBuf key)
+{
+    TPyObjectPtr kv = ConvertStringToPyObject(key);
+    if (!kv) {
+        PyErr_Clear();
+        return nullptr;
+    }
+    return PyDict_GetItem(v, kv.Get());
+}
+
+PyObject* GetItemFromPyDictUsingBytes(PyObject* v, TStringBuf key)
+{
+    TPyObjectPtr bytesMemberName = PyBytes_FromStringAndSize(key.data(), key.size());
+    return PyDict_GetItem(v, bytesMemberName.Get());
+}
+
+// Helper function for double-probing UTF-8 str() and bytes() in the case of Python 3.
+PyObject* GetItemFromPyDict(PyObject* v, TStringBuf key)
+{
+    PyObject* ret = GetItemFromPyDictUsingString(v, key);
+#if PY_MAJOR_VERSION >= 3
+    if (!ret) {
+        ret = GetItemFromPyDictUsingBytes(v, key);
+    }
+#endif
+    return ret;
+}
+
+// Sized counterpart of PyObject_GetAttrString from C API. Returns strong reference.
+TPyObjectPtr GetAttrFromPyObject(PyObject* v, TStringBuf name)
+{
+    TPyObjectPtr w = ConvertStringToPyObject(name);
+    if (!w) {
+        return nullptr;
+    }
+    return PyObject_GetAttr(v, w.Get());
+}
+
 }
 
 TPyObjectPtr ToPyStruct(const TPyCastContext::TPtr& ctx, const NUdf::TType* type, const NUdf::TUnboxedValuePod& value)
@@ -128,11 +176,7 @@ NUdf::TUnboxedValue FromPyStruct(const TPyCastContext::TPtr& ctx, const NUdf::TT
             TStringBuf memberName = inspector.GetMemberName(i);
             auto memberType = inspector.GetMemberType(i);
             // borrowed reference - no need to manage ownership
-            PyObject* item = PyDict_GetItemString(value, memberName.data());
-            if (!item) {
-                TPyObjectPtr bytesMemberName = PyBytes_FromStringAndSize(memberName.data(), memberName.size());
-                item = PyDict_GetItem(value, bytesMemberName.Get());
-            }
+            PyObject* item = GetItemFromPyDict(value, memberName);
             if (!item) {
                 if (ctx->PyCtx->TypeInfoHelper->GetTypeKind(memberType) == NUdf::ETypeKind::Optional) {
                     items[i] = NUdf::TUnboxedValue();
@@ -144,7 +188,7 @@ NUdf::TUnboxedValue FromPyStruct(const TPyCastContext::TPtr& ctx, const NUdf::TT
             }
 
             try {
-                items[i] = FromPyObject(ctx, inspector.GetMemberType(i), item);
+                items[i] = FromPyObject(ctx, memberType, item);
             } catch (const yexception& e) {
                 errors.push_back(TStringBuilder() << "Failed to convert dict item '" << memberName << "' - " << e.what());
             }
@@ -157,7 +201,7 @@ NUdf::TUnboxedValue FromPyStruct(const TPyCastContext::TPtr& ctx, const NUdf::TT
         for (ui32 i = 0; i < membersCount; i++) {
             TStringBuf memberName = inspector.GetMemberName(i);
             auto memberType = inspector.GetMemberType(i);
-            TPyObjectPtr attr = PyObject_GetAttrString(value, memberName.data());
+            TPyObjectPtr attr = GetAttrFromPyObject(value, memberName);
             if (!attr) {
                 if (ctx->PyCtx->TypeInfoHelper->GetTypeKind(memberType) == NUdf::ETypeKind::Optional &&
                     PyErr_ExceptionMatches(PyExc_AttributeError)) {