From aa419e61d9ebd8b535ea5b74c62eb92ffd63c6d9 Mon Sep 17 00:00:00 2001
From: Vadim Lebedev <vadim@mbdsys.com>
Date: Wed, 10 Aug 2022 00:50:43 +0200
Subject: [PATCH] String and unicode handling fixes

---
 libarchive/__init__.py        |  5 +----
 libarchive/_libarchive_wrap.c | 25 ++++++++++++++++++++++++-
 tests.py                      |  5 ++++-
 3 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/libarchive/__init__.py b/libarchive/__init__.py
index b95890f..df8c115 100644
--- a/libarchive/__init__.py
+++ b/libarchive/__init__.py
@@ -281,10 +281,7 @@ class EntryWriteStream(object):
         if self.closed:
             raise Exception('Cannot write to closed stream.')
         if self.buffer:
-            if PY3:
-                self.buffer.write(data)
-            elif isinstance(data, str):
-                self.buffer.write(unicode(data, ENCODING))
+            self.buffer.write(data)
         else:
             _libarchive.archive_write_data_from_str(self.archive._a, data.encode(ENCODING))
         self.bytes += len(data)
diff --git a/libarchive/_libarchive_wrap.c b/libarchive/_libarchive_wrap.c
index bcc604f..44a3147 100644
--- a/libarchive/_libarchive_wrap.c
+++ b/libarchive/_libarchive_wrap.c
@@ -3193,7 +3193,7 @@ SWIG_AsVal_unsigned_SS_short (PyObject * obj, unsigned short *val)
 }
 
 
-
+#if 1
 PyObject *archive_read_data_into_str(struct archive *archive, int len) {
     PyObject *str = NULL;
     if (!(str = PyBytes_FromStringAndSize(NULL, len))) {
@@ -3216,6 +3216,29 @@ PyObject *archive_write_data_from_str(struct archive *archive, PyObject *str) {
     }
     return PyInt_FromLong(len);
 }
+#else
+PyObject *archive_read_data_into_str(struct archive *archive, int len) {
+    PyObject *str = NULL;
+    if (!(str = PyUnicode_FromStringAndSize(NULL, len))) {
+        PyErr_SetString(PyExc_MemoryError, "could not allocate string.");
+        return NULL;
+    }
+    if (len != archive_read_data(archive, PyString_AS_STRING(str), len)) {
+        PyErr_SetString(PyExc_RuntimeError, "could not read requested data.");
+        return NULL;
+    }
+    return str;
+}
+
+PyObject *archive_write_data_from_str(struct archive *archive, PyObject *str) {
+    int len = PyString_Size(str);
+    if (!archive_write_data(archive, PyString_AS_STRING(str), len)) {
+        PyErr_SetString(PyExc_RuntimeError, "could not write requested data.");
+        return NULL;
+    }
+    return PyInt_FromLong(len);
+}
+#endif
 
 #ifdef __cplusplus
 extern "C" {
diff --git a/tests.py b/tests.py
index 5135362..6c6ad7c 100644
--- a/tests.py
+++ b/tests.py
@@ -282,7 +282,10 @@ class TestPasswordProtection(unittest.TestCase):
 
     def test_read_with_password(self):
         z = ZipFile(ZIPPATH, 'r', password='pwd')
-        self.assertEqual(z.read(ITEM_NAME), ITEM_CONTENT)
+        if PY3:
+            self.assertEqual(z.read(ITEM_NAME), bytes(ITEM_CONTENT, 'utf-8'))
+        else:
+            self.assertEqual(z.read(ITEM_NAME), ITEM_CONTENT)
         z.close()
 
     def test_read_without_password(self):