diff --git a/mwclient/image.py b/mwclient/image.py index 8cf176406a63c30f4058fc46c083ac6e42db9c60..e072a6836feb88d0c00e7f1b22711284ec137037 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)