From 3d731ba641f58e67ca5af9419ff6293163e2186b Mon Sep 17 00:00:00 2001
From: Daniel Grana <dangra@gmail.com>
Date: Fri, 30 Apr 2010 09:41:11 -0300
Subject: [PATCH] url_query_cleaner: add exclude and non-unique parameters
 support, also remove untested exception catching code and add missing tests

---
 scrapy/tests/test_utils_url.py | 22 +++++++++++++-----
 scrapy/utils/url.py            | 41 +++++++++++++++++-----------------
 2 files changed, 37 insertions(+), 26 deletions(-)

diff --git a/scrapy/tests/test_utils_url.py b/scrapy/tests/test_utils_url.py
index 05a22e8db..29e7cf8e0 100644
--- a/scrapy/tests/test_utils_url.py
+++ b/scrapy/tests/test_utils_url.py
@@ -136,12 +136,22 @@ class UrlUtilsTest(unittest.TestCase):
                          'http://rmc-offers.co.uk/productlist.asp?BCat=newvalue&CatID=60')
 
     def test_url_query_cleaner(self):
-        self.assertEqual(url_query_cleaner("product.html?id=200&foo=bar&name=wired", 'id'),
-                         'product.html?id=200')
-        self.assertEqual(url_query_cleaner("product.html?id=200&foo=bar&name=wired", ['id', 'name']),
-                         'product.html?id=200&name=wired')
-        self.assertEqual(url_query_cleaner("product.html?id=200&foo=bar&name=wired#id20", ['id', 'foo']),
-                         'product.html?id=200&foo=bar')
+        self.assertEqual('product.html?id=200',
+                url_query_cleaner("product.html?id=200&foo=bar&name=wired", ['id']))
+        self.assertEqual('product.html?id=200&name=wired',
+                url_query_cleaner("product.html?id=200&foo=bar&name=wired", ['id', 'name']))
+        self.assertEqual('product.html?id',
+                url_query_cleaner("product.html?id&other=3&novalue=", ['id']))
+        self.assertEqual('product.html?d=1&d=2&d=3',
+                url_query_cleaner("product.html?d=1&e=b&d=2&d=3&other=other", ['d'], unique=False))
+        self.assertEqual('product.html?id=200&foo=bar',
+                url_query_cleaner("product.html?id=200&foo=bar&name=wired#id20", ['id', 'foo']))
+        self.assertEqual('product.html?foo=bar&name=wired',
+                url_query_cleaner("product.html?id=200&foo=bar&name=wired", ['id'], remove=True))
+        self.assertEqual('product.html?name=wired',
+                url_query_cleaner("product.html?id=2&foo=bar&name=wired", ['id', 'foo'], remove=True))
+        self.assertEqual('product.html?foo=bar&name=wired',
+                url_query_cleaner("product.html?id=2&foo=bar&name=wired", ['id', 'footo'], remove=True))
 
     def test_canonicalize_url(self):
         # simplest case
diff --git a/scrapy/utils/url.py b/scrapy/utils/url.py
index 7ce27636d..02340deae 100644
--- a/scrapy/utils/url.py
+++ b/scrapy/utils/url.py
@@ -85,29 +85,30 @@ def url_query_parameter(url, parameter, default=None, keep_blank_values=0):
         keep_blank_values=keep_blank_values)
     return queryparams.get(parameter, [default])[0]
 
-def url_query_cleaner(url, parameterlist=(), sep='&', kvsep='='):
-    """Clean url arguments leaving only those passed in the parameterlist"""
-    try:
-        url = urlparse.urldefrag(url)[0]
-        base, query = url.split('?', 1)
-        parameters = [pair.split(kvsep, 1) for pair in query.split(sep)]
-    except:
-        base = url
-        query = ""
-        parameters = []
+def url_query_cleaner(url, parameterlist=(), sep='&', kvsep='=', remove=False, unique=True):
+    """Clean url arguments leaving only those passed in the parameterlist.
+
+    If remove is True, leave only those not in parameterlist.
+    If unique is False, do not remove duplicated keys
+    """
+    url = urlparse.urldefrag(url)[0]
+    base, _, query = url.partition('?')
+    parameters = [ksv.partition(kvsep) for ksv in query.split(sep)]
 
     # unique parameters while keeping order
-    unique = {}
+    seen = set()
     querylist = []
-    for pair in parameters:
-        k = pair[0]
-        if not unique.get(k):
-            querylist += [pair]
-            unique[k] = 1
-
-    query = sep.join([kvsep.join(pair) for pair in querylist if pair[0] in \
-        parameterlist])
-    return '?'.join([base, query])
+    for k, s, v in parameters:
+        if unique and k in seen:
+            continue
+        elif remove and k in parameterlist:
+            continue
+        elif not remove and k not in parameterlist:
+            continue
+        else:
+            querylist.append([k, s, v])
+            seen.add(k)
+    return base + '?' + sep.join(''.join(ksv) for ksv in querylist)
 
 def add_or_replace_parameter(url, name, new_value, sep='&', url_is_quoted=False):
     """Add or remove a parameter to a given url"""