From 038b22273325a2f57f542f6ee8501df3cbe8c597 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dan=20Michael=20O=2E=20Hegg=C3=B8?= <danmichaelo@gmail.com>
Date: Sun, 19 Jul 2015 03:32:14 +0200
Subject: [PATCH] Fix Image.download method
- The method has been broken since the move to requests (wrong number of args sent to self.site.connection.get)
- Add option to stream the file to a file object
---
mwclient/image.py | 26 ++++++++++++++++++++------
1 file changed, 20 insertions(+), 6 deletions(-)
diff --git a/mwclient/image.py b/mwclient/image.py
index 8cf1764..e072a68 100644
--- a/mwclient/image.py
+++ b/mwclient/image.py
@@ -25,13 +25,27 @@ class Image(mwclient.page.Page):
def duplicatefiles(self, limit=None):
return mwclient.listing.PageProperty(self, 'duplicatefiles', 'df', dflimit=limit)
- def download(self):
+ def download(self, destination=None):
+ """
+ Download the file. If `destination` is given, the file will be written
+ directly to the stream. Otherwise the file content will be stored in memory
+ and returned (with the risk of running out of memory for large files).
+
+ Recommended usage:
+
+ >>> with open(filename, 'wb') as fd:
+ ... image.download(fd)
+
+ Args:
+ destination (file object): Destination file
+ """
url = self.imageinfo['url']
- if not url.startswith('http://'):
- url = 'http://' + self.site.host + url
- url = urllib.parse.urlparse(url)
- # TODO: query string
- return self.site.connection.get(url[1], url[2])
+ if destination is not None:
+ res = self.site.connection.get(url, stream=True)
+ for chunk in res.iter_content(1024):
+ destination.write(chunk)
+ else:
+ return self.site.connection.get(url).content
def __repr__(self):
return "<Image object '%s' for %s>" % (self.name.encode('utf-8'), self.site)
--
GitLab