From 10ced29e18a4408c989c0ec39741fa3051cb3d37 Mon Sep 17 00:00:00 2001 From: Pablo Hoffman Date: Sun, 23 Oct 2011 02:49:17 -0200 Subject: [PATCH] changed feed exports storage api so that file/stdio outputs directly without using a temporary file --- scrapy/contrib/feedexport.py | 65 +++++++++++++++---------- scrapy/tests/test_contrib_feedexport.py | 46 +++++++++++------ 2 files changed, 68 insertions(+), 43 deletions(-) diff --git a/scrapy/contrib/feedexport.py b/scrapy/contrib/feedexport.py index a1e215367..28919b843 100644 --- a/scrapy/contrib/feedexport.py +++ b/scrapy/contrib/feedexport.py @@ -9,7 +9,6 @@ from tempfile import TemporaryFile from datetime import datetime from urlparse import urlparse from ftplib import FTP -from shutil import copyfileobj from zope.interface import Interface, implements from twisted.internet import defer, threads @@ -29,7 +28,11 @@ class IFeedStorage(Interface): def __init__(uri): """Initialize the storage with the parameters given in the URI""" - def store(file, spider): + def open(spider): + """Open the storage for the given spider. It must return a file-like + object that will be used for the exporters""" + + def store(file): """Store the given file stream""" @@ -37,10 +40,13 @@ class BlockingFeedStorage(object): implements(IFeedStorage) - def store(self, file, spider): - return threads.deferToThread(self._store_in_thread, file, spider) + def open(self, spider): + return TemporaryFile(prefix='feed-') - def _store_in_thread(self, file, spider): + def store(self, file): + return threads.deferToThread(self._store_in_thread, file) + + def _store_in_thread(self, file): raise NotImplementedError @@ -51,23 +57,27 @@ class StdoutFeedStorage(object): def __init__(self, uri, _stdout=sys.stdout): self._stdout = _stdout - def store(self, file, spider): - copyfileobj(file, self._stdout) + def open(self, spider): + return self._stdout + def store(self, file): + pass -class FileFeedStorage(BlockingFeedStorage): +class FileFeedStorage(object): + + implements(IFeedStorage) def __init__(self, uri): self.path = file_uri_to_path(uri) - def _store_in_thread(self, file, spider): + def open(self, spider): dirname = os.path.dirname(self.path) if dirname and not os.path.exists(dirname): os.makedirs(dirname) - f = open(self.path, 'wb') - copyfileobj(file, f) - f.close() + return open(self.path, 'ab') + def store(self, file): + file.close() class S3FeedStorage(BlockingFeedStorage): @@ -83,7 +93,8 @@ class S3FeedStorage(BlockingFeedStorage): self.secret_key = u.password or settings['AWS_SECRET_ACCESS_KEY'] self.keyname = u.path - def _store_in_thread(self, file, spider): + def _store_in_thread(self, file): + file.seek(0) conn = self.connect_s3(self.access_key, self.secret_key) bucket = conn.get_bucket(self.bucketname, validate=False) key = bucket.new_key(self.keyname) @@ -101,7 +112,8 @@ class FTPFeedStorage(BlockingFeedStorage): self.password = u.password self.path = u.path - def _store_in_thread(self, file, spider): + def _store_in_thread(self, file): + file.seek(0) ftp = FTP() ftp.connect(self.host, self.port) ftp.login(self.username, self.password) @@ -112,9 +124,11 @@ class FTPFeedStorage(BlockingFeedStorage): class SpiderSlot(object): - def __init__(self, file, exp): + def __init__(self, file, exporter, storage, uri): self.file = file - self.exporter = exp + self.exporter = exporter + self.storage = storage + self.uri = uri self.itemcount = 0 class FeedExporter(object): @@ -139,10 +153,12 @@ class FeedExporter(object): dispatcher.connect(self.item_scraped, signals.item_scraped) def open_spider(self, spider): - file = TemporaryFile(prefix='feed-') - exp = self._get_exporter(file) - exp.start_exporting() - self.slots[spider] = SpiderSlot(file, exp) + uri = self.urifmt % self._get_uri_params(spider) + storage = self._get_storage(uri) + file = storage.open(spider) + exporter = self._get_exporter(file) + exporter.start_exporting() + self.slots[spider] = SpiderSlot(file, exporter, storage, uri) def close_spider(self, spider): slot = self.slots.pop(spider) @@ -150,14 +166,9 @@ class FeedExporter(object): return slot.exporter.finish_exporting() nbytes = slot.file.tell() - slot.file.seek(0) - uri = self.urifmt % self._get_uri_params(spider) - storage = self._get_storage(uri) - if not storage: - return logfmt = "%%s %s feed (%d items, %d bytes) in: %s" % (self.format, \ - slot.itemcount, nbytes, uri) - d = defer.maybeDeferred(storage.store, slot.file, spider) + slot.itemcount, nbytes, slot.uri) + d = defer.maybeDeferred(slot.storage.store, slot.file) d.addCallback(lambda _: log.msg(logfmt % "Stored", spider=spider)) d.addErrback(log.err, logfmt % "Error storing", spider=spider) d.addBoth(lambda _: slot.file.close()) diff --git a/scrapy/tests/test_contrib_feedexport.py b/scrapy/tests/test_contrib_feedexport.py index 5e07b3f14..5c50eb601 100644 --- a/scrapy/tests/test_contrib_feedexport.py +++ b/scrapy/tests/test_contrib_feedexport.py @@ -10,18 +10,7 @@ from scrapy.spider import BaseSpider from scrapy.contrib.feedexport import IFeedStorage, FileFeedStorage, FTPFeedStorage, S3FeedStorage, StdoutFeedStorage from scrapy.utils.test import assert_aws_environ -class FeedStorageTest(unittest.TestCase): - - @defer.inlineCallbacks - def _assert_stores(self, storage, path): - yield storage.store(StringIO("content"), BaseSpider("default")) - self.failUnless(os.path.exists(path)) - self.failUnlessEqual(open(path).read(), "content") - # again, to check files are overwritten properly - yield storage.store(StringIO("new content"), BaseSpider("default")) - self.failUnlessEqual(open(path).read(), "new content") - -class FileFeedStorageTest(FeedStorageTest): +class FileFeedStorageTest(unittest.TestCase): def test_store_file_uri(self): path = os.path.abspath(self.mktemp()) @@ -47,8 +36,17 @@ class FileFeedStorageTest(FeedStorageTest): st = FileFeedStorage(path) verifyObject(IFeedStorage, st) + @defer.inlineCallbacks + def _assert_stores(self, storage, path): + spider = BaseSpider("default") + file = storage.open(spider) + file.write("content") + yield storage.store(file) + self.failUnless(os.path.exists(path)) + self.failUnlessEqual(open(path).read(), "content") -class FTPFeedStorageTest(FeedStorageTest): + +class FTPFeedStorageTest(unittest.TestCase): def test_store(self): uri = os.environ.get('FEEDTEST_FTP_URI') @@ -59,6 +57,18 @@ class FTPFeedStorageTest(FeedStorageTest): verifyObject(IFeedStorage, st) return self._assert_stores(st, path) + @defer.inlineCallbacks + def _assert_stores(self, storage, path): + spider = BaseSpider("default") + file = storage.open(spider) + file.write("content") + yield storage.store(file) + self.failUnless(os.path.exists(path)) + self.failUnlessEqual(open(path).read(), "content") + # again, to check s3 objects are overwritten + yield storage.store(StringIO("new content")) + self.failUnlessEqual(open(path).read(), "new content") + class S3FeedStorageTest(unittest.TestCase): @@ -71,16 +81,20 @@ class S3FeedStorageTest(unittest.TestCase): from boto import connect_s3 storage = S3FeedStorage(uri) verifyObject(IFeedStorage, storage) - yield storage.store(StringIO("content"), BaseSpider("default")) + file = storage.open(BaseSpider("default")) + file.write("content") + yield storage.store(file) u = urlparse.urlparse(uri) key = connect_s3().get_bucket(u.hostname, validate=False).get_key(u.path) self.failUnlessEqual(key.get_contents_as_string(), "content") -class StdoutFeedStorageTest(FeedStorageTest): +class StdoutFeedStorageTest(unittest.TestCase): @defer.inlineCallbacks def test_store(self): out = StringIO() storage = StdoutFeedStorage('stdout:', _stdout=out) - yield storage.store(StringIO("content"), BaseSpider("default")) + file = storage.open(BaseSpider("default")) + file.write("content") + yield storage.store(file) self.assertEqual(out.getvalue(), "content")