WPD: put_file() and add create_folder() method to drivers

src/calibre/devices/mtp/filesystem_cache.py

@@ -12,7 +12,7 @@ from operator import attrgetter
 from future_builtins import map
 
 from calibre import human_readable, prints, force_unicode
-from calibre.utils.icu import sort_key
+from calibre.utils.icu import sort_key, lower
 
 class FileOrFolder(object):
 
@@ -30,6 +30,9 @@ class FileOrFolder(object):
             if sid not in all_storage_ids:
                 sid = all_storage_ids[0]
             self.parent_id = sid
+        if self.parent_id is None and self.storage_id is None:
+            # A storage object
+            self.storage_id = self.object_id
         self.is_hidden = entry.get('is_hidden', False)
         self.is_system = entry.get('is_system', False)
         self.can_delete = entry.get('can_delete', True)
@@ -53,6 +56,12 @@ class FileOrFolder(object):
         for e in self.files:
             yield e
 
+    def add_child(self, entry):
+        ans = FileOrFolder(entry, self.id_map)
+        t = self.folders if ans.is_folder else self.files
+        t.append(ans)
+        return ans
+
     def dump(self, prefix='', out=sys.stdout):
         c = '+' if self.is_folder else '-'
         data = ('%s children'%(sum(map(len, (self.files, self.folders))))
@@ -63,6 +72,20 @@ class FileOrFolder(object):
             for e in sorted(c, key=lambda x:sort_key(x.name)):
                 e.dump(prefix=prefix+'  ', out=out)
 
+    def folder_named(self, name):
+        name = lower(name)
+        for e in self.folders:
+            if e.name and lower(e.name) == name:
+                return e
+        return None
+
+    def file_named(self, name):
+        name = lower(name)
+        for e in self.files:
+            if e.name and lower(e.name) == name:
+                return e
+        return None
+
 class FilesystemCache(object):
 
     def __init__(self, all_storage, entries):

src/calibre/devices/mtp/unix/driver.py

@@ -186,6 +186,14 @@ class MTP_DEVICE(MTPDeviceBase):
                 ans[i] = s['freespace_bytes']
         return tuple(ans)
 
+    @synchronous
+    def create_folder(self, parent_id, name):
+        parent = self.filesystem_cache.id_map[parent_id]
+        e = parent.folder_named(name)
+        if e is not None:
+            return e
+        ans = self.dev.create_folder(parent.storage_id, parent_id, name)
+        return parent.add_child(ans)
 
 if __name__ == '__main__':
     BytesIO

src/calibre/devices/mtp/windows/content_enumeration.cpp

@@ -372,7 +372,7 @@ end:
 } 
 // }}}
 
-static IPortableDeviceValues* create_object_properties(const wchar_t *parent_id, const wchar_t *name, const GUID content_type) { // {{{
+static IPortableDeviceValues* create_object_properties(const wchar_t *parent_id, const wchar_t *name, const GUID content_type, unsigned PY_LONG_LONG size) { // {{{
     IPortableDeviceValues *values = NULL;
     HRESULT hr;
     BOOL ok = FALSE;
@@ -396,6 +396,11 @@ static IPortableDeviceValues* create_object_properties(const wchar_t *parent_id,
     hr = values->SetGuidValue(WPD_OBJECT_CONTENT_TYPE, content_type);
     if (FAILED(hr)) { hresult_set_exc("Failed to set content_type value", hr); goto end; }
 
+    if (!IsEqualGUID(WPD_CONTENT_TYPE_FOLDER, content_type)) {
+        hr = values->SetUnsignedLargeIntegerValue(WPD_OBJECT_SIZE, size);
+        if (FAILED(hr)) { hresult_set_exc("Failed to set size value", hr); goto end; }
+    }
+
     ok = TRUE;
 
 end:
@@ -538,11 +543,11 @@ PyObject* wpd::create_folder(IPortableDevice *device, const wchar_t *parent_id,
     IPortableDeviceValues *values = NULL;
     IPortableDeviceProperties *devprops = NULL;
     IPortableDeviceKeyCollection *properties = NULL;
+    wchar_t *newid = NULL;
     PyObject *ans = NULL;
     HRESULT hr;
-    wchar_t *newid = NULL;
 
-    values = create_object_properties(parent_id, name, WPD_CONTENT_TYPE_FOLDER);
+    values = create_object_properties(parent_id, name, WPD_CONTENT_TYPE_FOLDER, 0);
     if (values == NULL) goto end;
 
     Py_BEGIN_ALLOW_THREADS;
@@ -612,4 +617,91 @@ end:
 
 } // }}}
 
+PyObject* wpd::put_file(IPortableDevice *device, const wchar_t *parent_id, const wchar_t *name, PyObject *src, unsigned PY_LONG_LONG size, PyObject *callback) { // {{{
+    IPortableDeviceContent *content = NULL;
+    IPortableDeviceValues *values = NULL;
+    IPortableDeviceProperties *devprops = NULL;
+    IPortableDeviceKeyCollection *properties = NULL;
+    IStream *temp = NULL;
+    IPortableDeviceDataStream *dest = NULL;
+    char *buf = NULL;
+    wchar_t *newid = NULL;
+    PyObject *ans = NULL, *raw;
+    HRESULT hr;
+    DWORD bufsize = 0;
+    BOOL ok = FALSE;
+    Py_ssize_t bytes_read = 0;
+    ULONG bytes_written = 0, total_written = 0;
+
+    values = create_object_properties(parent_id, name, WPD_CONTENT_TYPE_GENERIC_FILE, size);
+    if (values == NULL) goto end;
+
+    Py_BEGIN_ALLOW_THREADS;
+    hr = device->Content(&content);
+    Py_END_ALLOW_THREADS;
+    if (FAILED(hr)) { hresult_set_exc("Failed to create content interface", hr); goto end; }
+
+    hr = content->Properties(&devprops);
+    if (FAILED(hr)) { hresult_set_exc("Failed to get IPortableDeviceProperties interface", hr); goto end; }
+
+    properties = create_filesystem_properties_collection();
+    if (properties == NULL) goto end;
+
+    Py_BEGIN_ALLOW_THREADS;
+    hr = content->CreateObjectWithPropertiesAndData(values, &temp, &bufsize, NULL);
+    Py_END_ALLOW_THREADS;
+    if (FAILED(hr)) { 
+        if (HRESULT_FROM_WIN32(ERROR_BUSY) == hr) {
+            PyErr_SetString(WPDFileBusy, "Object is in use");
+        } else hresult_set_exc("Failed to create stream interface to write to object", hr); 
+        goto end; 
+    }
+
+    hr = temp->QueryInterface(IID_PPV_ARGS(&dest));
+    if (FAILED(hr)) { hresult_set_exc("Failed to create IPortableDeviceStream", hr); goto end; }
+
+    while(TRUE) {
+        raw = PyObject_CallMethod(src, "read", "k", bufsize);
+        if (raw == NULL) break;
+        PyBytes_AsStringAndSize(raw, &buf, &bytes_read);
+        if (bytes_read > 0) {
+            Py_BEGIN_ALLOW_THREADS;
+            hr = dest->Write(buf, bytes_read, &bytes_written);
+            Py_END_ALLOW_THREADS;
+            Py_DECREF(raw);
+            if (hr == STG_E_MEDIUMFULL) { PyErr_SetString(WPDError, "Cannot write to device as it is full"); break; }
+            if (hr == STG_E_ACCESSDENIED) { PyErr_SetString(WPDError, "Cannot write to file as access is denied"); break; }
+            if (hr == STG_E_WRITEFAULT) { PyErr_SetString(WPDError, "Cannot write to file as there was a disk I/O error"); break; }
+            if (FAILED(hr)) { hresult_set_exc("Cannot write to file", hr); break; }
+            if (bytes_written != bytes_read) { PyErr_SetString(WPDError, "Writing to file failed, not all bytes were written"); break; }
+            total_written += bytes_written;
+            if (callback != NULL) Py_XDECREF(PyObject_CallFunction(callback, "kK", total_written, size));
+        } else Py_DECREF(raw);
+        if (bytes_read == 0) { ok = TRUE; break; }
+    }
+    if (!ok) {dest->Revert(); goto end;}
+    Py_BEGIN_ALLOW_THREADS;
+    hr = dest->Commit(STGC_DEFAULT);
+    Py_END_ALLOW_THREADS;
+    if (FAILED(hr)) { hresult_set_exc("Failed to write data to file, commit failed", hr); goto end; }
+    if (callback != NULL) Py_XDECREF(PyObject_CallFunction(callback, "kK", total_written, size));
+
+    Py_BEGIN_ALLOW_THREADS;
+    hr = dest->GetObjectID(&newid);
+    Py_END_ALLOW_THREADS;
+    if (FAILED(hr)) { hresult_set_exc("Failed to get id of newly created file", hr); goto end; }
+
+    ans = get_object_properties(devprops, properties, newid);
+end:
+    if (content != NULL) content->Release();
+    if (values != NULL) values->Release();
+    if (devprops != NULL) devprops->Release();
+    if (properties != NULL) properties->Release();
+    if (temp != NULL) temp->Release();
+    if (dest != NULL) dest->Release();
+    if (newid != NULL) CoTaskMemFree(newid);
+    return ans;
+
+} // }}}
+
 } // namespace wpd

src/calibre/devices/mtp/windows/device.cpp

@@ -138,6 +138,25 @@ py_delete_object(Device *self, PyObject *args, PyObject *kwargs) {
     return ret;
 } // }}}
 
+// get_file() {{{
+static PyObject*
+py_put_file(Device *self, PyObject *args, PyObject *kwargs) {
+    PyObject *pparent_id, *pname, *stream, *callback = NULL, *ret;
+    wchar_t *parent_id, *name;
+    unsigned PY_LONG_LONG size;
+
+    if (!PyArg_ParseTuple(args, "OOOK|O", &pparent_id, &pname, &stream, &size, &callback)) return NULL;
+    parent_id = unicode_to_wchar(pparent_id);
+    name = unicode_to_wchar(pname);
+    if (parent_id == NULL || name == NULL) return NULL;
+
+    if (callback == NULL || !PyCallable_Check(callback)) callback = NULL;
+
+    ret = wpd::put_file(self->device, parent_id, name, stream, size, callback);
+    free(parent_id); free(name);
+    return ret;
+} // }}}
+
 static PyMethodDef Device_methods[] = {
     {"update_data", (PyCFunction)update_data, METH_VARARGS,
      "update_data() -> Reread the basic device data from the device (total, space, free space, storage locations, etc.)"
@@ -159,6 +178,10 @@ static PyMethodDef Device_methods[] = {
      "delete_object(object_id) -> Delete the object identified by object_id. Note that trying to delete a non-empty folder will raise an error."
     },
 
+    {"put_file", (PyCFunction)py_put_file, METH_VARARGS,
+     "put_file(parent_id, name, stream, size_in_bytes, callback=None) -> Copy a file from the stream object, creating a new file on the device with parent identified by parent_id. Returns the file metadata of the newly created file. callback should be a callable that accepts two argument: (bytes_written, total_size). It will be called after each chunk is written to the device. Note that it can be called multiple times with the same arguments."
+    },
+
     {NULL}
 };
 

src/calibre/devices/mtp/windows/driver.py

@@ -254,3 +254,13 @@ class MTP_DEVICE(MTPDeviceBase):
             self.dev.get_file(object_id, stream, callback)
         return stream
 
+    @same_thread
+    def create_folder(self, parent_id, name):
+        parent = self.filesystem_cache.id_map[parent_id]
+        e = parent.folder_named(name)
+        if e is not None:
+            return e
+        ans = self.dev.create_folder(parent_id, name)
+        ans['storage_id'] = parent.storage_id
+        return parent.add_child(ans)
+

src/calibre/devices/mtp/windows/global.h

@@ -60,6 +60,7 @@ extern PyObject* get_filesystem(IPortableDevice *device, const wchar_t *storage_
 extern PyObject* get_file(IPortableDevice *device, const wchar_t *object_id, PyObject *dest, PyObject *callback);
 extern PyObject* create_folder(IPortableDevice *device, const wchar_t *parent_id, const wchar_t *name);
 extern PyObject* delete_object(IPortableDevice *device, const wchar_t *object_id);
+extern PyObject* put_file(IPortableDevice *device, const wchar_t *parent_id, const wchar_t *name, PyObject *src, unsigned PY_LONG_LONG size, PyObject *callback);
 
 }