From e81120e68356cb0d880ec839fd42705b6bf1f4e5 Mon Sep 17 00:00:00 2001 From: Martin Olveyra Date: Wed, 2 Oct 2013 18:23:13 -0200 Subject: [PATCH] allow to set source and destination fields in files and images pipelines, and added tests --- scrapy/contrib/pipeline/files.py | 10 +++++--- scrapy/contrib/pipeline/images.py | 11 ++++++--- scrapy/tests/test_pipeline_files.py | 36 +++++++++++++++++++++++++++- scrapy/tests/test_pipeline_images.py | 34 ++++++++++++++++++++++++++ 4 files changed, 84 insertions(+), 7 deletions(-) diff --git a/scrapy/contrib/pipeline/files.py b/scrapy/contrib/pipeline/files.py index f03bf69df..1b2fef57d 100644 --- a/scrapy/contrib/pipeline/files.py +++ b/scrapy/contrib/pipeline/files.py @@ -141,6 +141,8 @@ class FilesPipeline(MediaPipeline): 'file': FSFilesStore, 's3': S3FilesStore, } + DEFAULT_FILES_URLS_FIELD = 'file_urls' + DEFAULT_FILES_RESULT_FIELD = 'files' def __init__(self, store_uri, download_func=None): if not store_uri: @@ -154,6 +156,8 @@ class FilesPipeline(MediaPipeline): s3store.AWS_ACCESS_KEY_ID = settings['AWS_ACCESS_KEY_ID'] s3store.AWS_SECRET_ACCESS_KEY = settings['AWS_SECRET_ACCESS_KEY'] + cls.FILES_URLS_FIELD = settings.get('FILES_URLS_FIELD', cls.DEFAULT_FILES_URLS_FIELD) + cls.FILES_RESULT_FIELD = settings.get('FILES_RESULT_FIELD', cls.DEFAULT_FILES_RESULT_FIELD) cls.EXPIRES = settings.getint('FILES_EXPIRES', 90) store_uri = settings['FILES_STORE'] return cls(store_uri) @@ -248,7 +252,7 @@ class FilesPipeline(MediaPipeline): ### Overridable Interface def get_media_requests(self, item, info): - return [Request(x) for x in item.get('file_urls', [])] + return [Request(x) for x in item.get(self.FILES_URLS_FIELD, [])] def file_key(self, url): media_guid = hashlib.sha1(url).hexdigest() @@ -263,6 +267,6 @@ class FilesPipeline(MediaPipeline): return checksum def item_completed(self, results, item, info): - if 'files' in item.fields: - item['files'] = [x for ok, x in results if ok] + if self.FILES_RESULT_FIELD in item.fields: + item[self.FILES_RESULT_FIELD] = [x for ok, x in results if ok] return item diff --git a/scrapy/contrib/pipeline/images.py b/scrapy/contrib/pipeline/images.py index aaad689a0..73e4998c5 100644 --- a/scrapy/contrib/pipeline/images.py +++ b/scrapy/contrib/pipeline/images.py @@ -33,6 +33,8 @@ class ImagesPipeline(FilesPipeline): MIN_WIDTH = 0 MIN_HEIGHT = 0 THUMBS = {} + DEFAULT_IMAGES_URLS_FIELD = 'image_urls' + DEFAULT_IMAGES_RESULT_FIELD = 'images' @classmethod def from_settings(cls, settings): @@ -43,6 +45,9 @@ class ImagesPipeline(FilesPipeline): s3store = cls.STORE_SCHEMES['s3'] s3store.AWS_ACCESS_KEY_ID = settings['AWS_ACCESS_KEY_ID'] s3store.AWS_SECRET_ACCESS_KEY = settings['AWS_SECRET_ACCESS_KEY'] + + cls.IMAGES_URLS_FIELD = settings.get('IMAGES_URLS_FIELD', cls.DEFAULT_IMAGES_URLS_FIELD) + cls.IMAGES_RESULT_FIELD = settings.get('IMAGES_RESULT_FIELD', cls.DEFAULT_IMAGES_RESULT_FIELD) store_uri = settings['IMAGES_STORE'] return cls(store_uri) @@ -103,7 +108,7 @@ class ImagesPipeline(FilesPipeline): return 'thumbs/%s/%s.jpg' % (thumb_id, image_guid) def get_media_requests(self, item, info): - return [Request(x) for x in item.get('image_urls', [])] + return [Request(x) for x in item.get(self.IMAGES_URLS_FIELD, [])] # backwards compatibility def image_key(self, url): @@ -111,6 +116,6 @@ class ImagesPipeline(FilesPipeline): return 'full/%s.jpg' % (media_guid) def item_completed(self, results, item, info): - if 'images' in item.fields: - item['images'] = [x for ok, x in results if ok] + if self.IMAGES_RESULT_FIELD in item.fields: + item[self.IMAGES_RESULT_FIELD] = [x for ok, x in results if ok] return item diff --git a/scrapy/tests/test_pipeline_files.py b/scrapy/tests/test_pipeline_files.py index f0aadfbfa..d9d20a658 100644 --- a/scrapy/tests/test_pipeline_files.py +++ b/scrapy/tests/test_pipeline_files.py @@ -10,6 +10,7 @@ from twisted.internet import defer from scrapy.contrib.pipeline.files import FilesPipeline, FSFilesStore from scrapy.item import Item, Field from scrapy.http import Request, Response +from scrapy.settings import Settings def _mocked_download_func(request, info): @@ -21,7 +22,8 @@ class FilesPipelineTestCase(unittest.TestCase): def setUp(self): self.tempdir = mkdtemp() - self.pipeline = FilesPipeline(self.tempdir, download_func=_mocked_download_func) + self.pipeline = FilesPipeline.from_settings(Settings({'FILES_STORE': self.tempdir})) + self.pipeline.download_func = _mocked_download_func self.pipeline.open_spider(None) def tearDown(self): @@ -87,7 +89,39 @@ class FilesPipelineTestCase(unittest.TestCase): map(lambda p: p.stop(), patchers) +class FilesPipelineTestCaseFields(unittest.TestCase): + def test_item_fields_default(self): + from scrapy.contrib.pipeline.files import FilesPipeline + class TestItem(Item): + name = Field() + file_urls = Field() + files = Field() + url = 'http://www.example.com/files/1.txt' + item = TestItem({'name': 'item1', 'file_urls': [url]}) + pipeline = FilesPipeline.from_settings(Settings({'FILES_STORE': 's3://example/files/'})) + requests = list(pipeline.get_media_requests(item, None)) + self.assertEqual(requests[0].url, url) + results = [(True, {'url': url})] + pipeline.item_completed(results, item, None) + self.assertEqual(item['files'], [results[0][1]]) + + def test_item_fields_override_settings(self): + from scrapy.contrib.pipeline.files import FilesPipeline + class TestItem(Item): + name = Field() + files = Field() + stored_file = Field() + url = 'http://www.example.com/files/1.txt' + item = TestItem({'name': 'item1', 'files': [url]}) + pipeline = FilesPipeline.from_settings(Settings({'FILES_STORE': 's3://example/files/', + 'FILES_URLS_FIELD': 'files', 'FILES_RESULT_FIELD': 'stored_file'})) + requests = list(pipeline.get_media_requests(item, None)) + self.assertEqual(requests[0].url, url) + results = [(True, {'url': url})] + pipeline.item_completed(results, item, None) + self.assertEqual(item['stored_file'], [results[0][1]]) + class ItemWithFiles(Item): file_urls = Field() files = Field() diff --git a/scrapy/tests/test_pipeline_images.py b/scrapy/tests/test_pipeline_images.py index 692e91afa..ae607f3fb 100644 --- a/scrapy/tests/test_pipeline_images.py +++ b/scrapy/tests/test_pipeline_images.py @@ -5,6 +5,8 @@ from shutil import rmtree from twisted.trial import unittest +from scrapy.item import Item, Field +from scrapy.settings import Settings skip = False try: @@ -82,7 +84,39 @@ class ImagesPipelineTestCase(unittest.TestCase): self.assertEquals(converted.mode, 'RGB') self.assertEquals(converted.getcolors(), [(10000, (205, 230, 255))]) +class ImagesPipelineTestCaseFields(unittest.TestCase): + def test_item_fields_default(self): + from scrapy.contrib.pipeline.images import ImagesPipeline + class TestItem(Item): + name = Field() + image_urls = Field() + images = Field() + url = 'http://www.example.com/images/1.jpg' + item = TestItem({'name': 'item1', 'image_urls': [url]}) + pipeline = ImagesPipeline.from_settings(Settings({'IMAGES_STORE': 's3://example/images/'})) + requests = list(pipeline.get_media_requests(item, None)) + self.assertEqual(requests[0].url, url) + results = [(True, {'url': url})] + pipeline.item_completed(results, item, None) + self.assertEqual(item['images'], [results[0][1]]) + + def test_item_fields_override_settings(self): + from scrapy.contrib.pipeline.images import ImagesPipeline + class TestItem(Item): + name = Field() + image = Field() + stored_image = Field() + url = 'http://www.example.com/images/1.jpg' + item = TestItem({'name': 'item1', 'image': [url]}) + pipeline = ImagesPipeline.from_settings(Settings({'IMAGES_STORE': 's3://example/images/', + 'IMAGES_URLS_FIELD': 'image', 'IMAGES_RESULT_FIELD': 'stored_image'})) + requests = list(pipeline.get_media_requests(item, None)) + self.assertEqual(requests[0].url, url) + results = [(True, {'url': url})] + pipeline.item_completed(results, item, None) + self.assertEqual(item['stored_image'], [results[0][1]]) + def _create_image(format, *a, **kw): buf = StringIO() Image.new(*a, **kw).save(buf, format)