_sort/_next links now use new path_with_replaced_args method

datasette/templates/_rows_and_columns.html

@@ -7,9 +7,9 @@
                         {{ column.name }}
                     {% else %}
                         {% if column.name == sort %}
-                            <a href="{{ path_with_added_args(request, {'_sort_desc': column.name, '_sort': None, '_next': None}) }}" rel="nofollow">{{ column.name }}&nbsp;▼</a>
+                            <a href="{{ path_with_replaced_args(request, {'_sort_desc': column.name, '_sort': None, '_next': None}) }}" rel="nofollow">{{ column.name }}&nbsp;▼</a>
                         {% else %}
-                            <a href="{{ path_with_added_args(request, {'_sort': column.name, '_sort_desc': None, '_next': None}) }}" rel="nofollow">{{ column.name }}{% if column.name == sort_desc %}&nbsp;▲{% endif %}</a>
+                            <a href="{{ path_with_replaced_args(request, {'_sort': column.name, '_sort_desc': None, '_next': None}) }}" rel="nofollow">{{ column.name }}{% if column.name == sort_desc %}&nbsp;▲{% endif %}</a>
                         {% endif %}
                     {% endif %}
                 </th>

datasette/utils.py

@@ -177,6 +177,22 @@ def path_with_removed_args(request, args, path=None):
     return path + query_string
 
 
+def path_with_replaced_args(request, args, path=None):
+    path = path or request.path
+    if isinstance(args, dict):
+        args = args.items()
+    keys_to_replace = {p[0] for p in args}
+    current = []
+    for key, value in urllib.parse.parse_qsl(request.query_string):
+        if key not in keys_to_replace:
+            current.append((key, value))
+    current.extend([p for p in args if p[1] is not None])
+    query_string = urllib.parse.urlencode(current)
+    if query_string:
+        query_string = '?{}'.format(query_string)
+    return path + query_string
+
+
 def path_with_ext(request, ext):
     path = request.path
     path += ext

datasette/views/table.py

@@ -14,6 +14,7 @@ from datasette.utils import (
     path_from_row_pks,
     path_with_added_args,
     path_with_removed_args,
+    path_with_replaced_args,
     to_css_class,
     urlsafe_components
 )
@@ -562,7 +563,7 @@ class TableView(RowTableShared):
             else:
                 added_args = {"_next": next_value}
             next_url = urllib.parse.urljoin(
-                request.url, path_with_added_args(request, added_args)
+                request.url, path_with_replaced_args(request, added_args)
             )
             rows = rows[:page_size]
 
@@ -650,7 +651,7 @@ class TableView(RowTableShared):
                 "filter_columns": filter_columns,
                 "display_rows": display_rows,
                 "is_sortable": any(c["sortable"] for c in display_columns),
-                "path_with_added_args": path_with_added_args,
+                "path_with_replaced_args": path_with_replaced_args,
                 "request": request,
                 "sort": sort,
                 "sort_desc": sort_desc,

tests/test_html.py

@@ -150,6 +150,57 @@ def test_sort_by_desc_redirects(app_client):
     assert response.headers['Location'].endswith('?_sort_desc=sortable')
 
 
+def test_sort_links(app_client):
+    response = app_client.get(
+        '/test_tables/sortable?_sort=sortable',
+        gather_request=False
+    )
+    assert response.status == 200
+    ths = Soup(response.body, 'html.parser').findAll('th')
+    attrs_and_link_attrs = [{
+        'attrs': th.attrs,
+        'a_href': (
+            th.find('a')['href'].split('/')[-1]
+            if th.find('a')
+            else None
+        ),
+    } for th in ths]
+    assert [
+        {
+            "attrs": {"class": ["col-Link"], "scope": "col"},
+            "a_href": None
+        },
+        {
+            "attrs": {"class": ["col-pk1"], "scope": "col"},
+            "a_href": None
+        },
+        {
+            "attrs": {"class": ["col-pk2"], "scope": "col"},
+            "a_href": None
+        },
+        {
+            "attrs": {"class": ["col-content"], "scope": "col"},
+            "a_href": None
+        },
+        {
+            "attrs": {"class": ["col-sortable"], "scope": "col"},
+            "a_href": "sortable?_sort_desc=sortable",
+        },
+        {
+            "attrs": {"class": ["col-sortable_with_nulls"], "scope": "col"},
+            "a_href": "sortable?_sort=sortable_with_nulls",
+        },
+        {
+            "attrs": {"class": ["col-sortable_with_nulls_2"], "scope": "col"},
+            "a_href": "sortable?_sort=sortable_with_nulls_2",
+        },
+        {
+            "attrs": {"class": ["col-text"], "scope": "col"},
+            "a_href": "sortable?_sort=text",
+        },
+    ] == attrs_and_link_attrs
+
+
 @pytest.mark.parametrize('path,expected_classes', [
     ('/', ['index']),
     ('/test_tables', ['db', 'db-test_tables']),

tests/test_utils.py

@@ -61,6 +61,19 @@ def test_path_with_removed_args(path, args, expected):
     assert expected == actual
 
 
+@pytest.mark.parametrize('path,args,expected', [
+    ('/foo?bar=1', {'bar': 2}, '/foo?bar=2'),
+    ('/foo?bar=1&baz=2', {'bar': None}, '/foo?baz=2'),
+])
+def test_path_with_replaced_args(path, args, expected):
+    request = Request(
+        path.encode('utf8'),
+        {}, '1.1', 'GET', None
+    )
+    actual = utils.path_with_replaced_args(request, args)
+    assert expected == actual
+
+
 @pytest.mark.parametrize('row,pks,expected_path', [
     ({'A': 'foo', 'B': 'bar'}, ['A', 'B'], 'foo,bar'),
     ({'A': 'f,o', 'B': 'bar'}, ['A', 'B'], 'f%2Co,bar'),