Intermediate changes

commit_hash:0897b988c8ca7235b4f13624177866ac14a46fdd

2 files changed, 74 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)) {
diff --git a/yql/essentials/udfs/common/python/bindings/py_struct_ut.cpp b/yql/essentials/udfs/common/python/bindings/py_struct_ut.cpp
index 9b232d99a3d..b8662907907 100644
--- a/yql/essentials/udfs/common/python/bindings/py_struct_ut.cpp
+++ b/yql/essentials/udfs/common/python/bindings/py_struct_ut.cpp
@@ -60,6 +60,29 @@ Y_UNIT_TEST_SUITE(TPyStructTest) {
                 });
     }
 
+    Y_UNIT_TEST(FromPyObjectBytesAttrWithNullCharacter) {
+        TPythonTestEngine engine;
+
+        ui32 ageIdx = 0;
+        auto personType = engine.GetTypeBuilder().Struct()->
+                AddField<int>("a\0ge", &ageIdx)
+                .Build();
+
+        engine.ToMiniKQL(personType,
+                "class Person:\n"
+                "    def __init__(self, age):\n"
+                "        setattr(self, 'a\\0ge', age)\n"
+                "\n"
+                "def Test():\n"
+                "    return Person(99)\n",
+                [ageIdx](const NUdf::TUnboxedValuePod& value) {
+                    UNIT_ASSERT(value);
+                    UNIT_ASSERT(value.IsBoxed());
+                    auto age = value.GetElement(ageIdx);
+                    UNIT_ASSERT_EQUAL(age.Get<ui32>(), 99);
+                });
+    }
+
     Y_UNIT_TEST(FromPyDict) {
         TPythonTestEngine engine;