From 10ed28b9d02ddc0b216889c7d37d452ac4b11324 Mon Sep 17 00:00:00 2001
From: Pablo Hoffman <pablo@pablohoffman.com>
Date: Tue, 3 Jan 2012 12:17:17 -0200
Subject: [PATCH] SitemapSpider: added support for sitemap urls ending in .xml
 and .xml.gz, even if they have a wrong content type

---
 scrapy/contrib/spiders/sitemap.py | 26 ++++++++++++++++----------
 scrapy/tests/test_spider.py       | 31 ++++++++++++++++++++++++++++++-
 scrapy/utils/gz.py                |  5 +++++
 3 files changed, 51 insertions(+), 11 deletions(-)

diff --git a/scrapy/contrib/spiders/sitemap.py b/scrapy/contrib/spiders/sitemap.py
index 698e7c389..91ae78c62 100644
--- a/scrapy/contrib/spiders/sitemap.py
+++ b/scrapy/contrib/spiders/sitemap.py
@@ -3,7 +3,7 @@ import re
 from scrapy.spider import BaseSpider
 from scrapy.http import Request, XmlResponse
 from scrapy.utils.sitemap import Sitemap, sitemap_urls_from_robots
-from scrapy.utils.gz import gunzip
+from scrapy.utils.gz import gunzip, is_gzipped
 from scrapy import log
 
 class SitemapSpider(BaseSpider):
@@ -29,12 +29,9 @@ class SitemapSpider(BaseSpider):
             for url in sitemap_urls_from_robots(response.body):
                 yield Request(url, callback=self._parse_sitemap)
         else:
-            if isinstance(response, XmlResponse):
-                body = response.body
-            elif is_gzipped(response):
-                body = gunzip(response.body)
-            else:
-                log.msg("Ignoring non-XML sitemap: %s" % response, log.WARNING)
+            body = self._get_sitemap_body(response)
+            if body is None:
+                log.msg("Ignoring invalid sitemap: %s" % response, log.WARNING)
                 return
 
             s = Sitemap(body)
@@ -49,9 +46,18 @@ class SitemapSpider(BaseSpider):
                             yield Request(loc, callback=c)
                             break
 
-def is_gzipped(response):
-    ctype = response.headers.get('Content-Type', '')
-    return ctype in ('application/x-gzip', 'application/gzip')
+    def _get_sitemap_body(self, response):
+        """Return the sitemap body contained in the given response, or None if the
+        response is not a sitemap.
+        """
+        if isinstance(response, XmlResponse):
+            return response.body
+        elif is_gzipped(response):
+            return gunzip(response.body)
+        elif response.url.endswith('.xml'):
+            return response.body
+        elif response.url.endswith('.xml.gz'):
+            return gunzip(response.body)
 
 def regex(x):
     if isinstance(x, basestring):
diff --git a/scrapy/tests/test_spider.py b/scrapy/tests/test_spider.py
index 4ab140311..11476ac65 100644
--- a/scrapy/tests/test_spider.py
+++ b/scrapy/tests/test_spider.py
@@ -1,10 +1,12 @@
 from __future__ import with_statement
 
-import warnings
+import gzip, warnings
+from cStringIO import StringIO
 
 from twisted.trial import unittest
 
 from scrapy.spider import BaseSpider
+from scrapy.http import Response, TextResponse, XmlResponse, HtmlResponse
 from scrapy.contrib.spiders.init import InitSpider
 from scrapy.contrib.spiders import CrawlSpider, XMLFeedSpider, CSVFeedSpider, SitemapSpider
 
@@ -55,6 +57,33 @@ class SitemapSpiderTest(BaseSpiderTest):
 
     spider_class = SitemapSpider
 
+    BODY = "SITEMAP"
+    f = StringIO()
+    g = gzip.GzipFile(fileobj=f, mode='w+b')
+    g.write(BODY)
+    g.close()
+    GZBODY = f.getvalue()
+
+    def test_get_sitemap_body(self):
+        spider = self.spider_class("example.com")
+
+        r = XmlResponse(url="http://www.example.com/", body=self.BODY)
+        self.assertEqual(spider._get_sitemap_body(r), self.BODY)
+
+        r = HtmlResponse(url="http://www.example.com/", body=self.BODY)
+        self.assertEqual(spider._get_sitemap_body(r), None)
+
+        r = Response(url="http://www.example.com/favicon.ico", body=self.BODY)
+        self.assertEqual(spider._get_sitemap_body(r), None)
+
+        r = Response(url="http://www.example.com/sitemap", body=self.GZBODY, headers={"content-type": "application/gzip"})
+        self.assertEqual(spider._get_sitemap_body(r), self.BODY)
+
+        r = TextResponse(url="http://www.example.com/sitemap.xml", body=self.BODY)
+        self.assertEqual(spider._get_sitemap_body(r), self.BODY)
+
+        r = Response(url="http://www.example.com/sitemap.xml.gz", body=self.GZBODY)
+        self.assertEqual(spider._get_sitemap_body(r), self.BODY)
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/scrapy/utils/gz.py b/scrapy/utils/gz.py
index 2dc9cd3d2..d6221da0d 100644
--- a/scrapy/utils/gz.py
+++ b/scrapy/utils/gz.py
@@ -21,3 +21,8 @@ def gunzip(data):
             else:
                 raise
     return output
+
+def is_gzipped(response):
+    """Return True if the response is gzipped, or False otherwise"""
+    ctype = response.headers.get('Content-Type', '')
+    return ctype in ('application/x-gzip', 'application/gzip')