From 700d83d8adfeb3859ebc93828951e5048cb0e425 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Mon, 23 Jul 2018 20:07:57 -0700
Subject: [PATCH] ?_json_infinity=1 for handling Infinity/-Infinity - fixes
 #332

---
 datasette/utils.py      | 12 ++++++++++++
 datasette/views/base.py | 10 +++++++++-
 docs/json_api.rst       |  6 ++++++
 tests/fixtures.py       |  9 +++++++++
 tests/test_api.py       | 29 ++++++++++++++++++++++++++++-
 5 files changed, 64 insertions(+), 2 deletions(-)

diff --git a/datasette/utils.py b/datasette/utils.py
index f1b403c0..7419f9ae 100644
--- a/datasette/utils.py
+++ b/datasette/utils.py
@@ -864,3 +864,15 @@ class LimitedWriter:
                 self.limit_bytes
             ))
         self.writer.write(bytes)
+
+
+_infinities = {float("inf"), float("-inf")}
+
+
+def remove_infinites(row):
+    if any((c in _infinities) if isinstance(c, float) else 0 for c in row):
+        return [
+            None if (isinstance(c, float) and c in _infinities) else c
+            for c in row
+        ]
+    return row
diff --git a/datasette/views/base.py b/datasette/views/base.py
index 44381cb9..3da3793c 100644
--- a/datasette/views/base.py
+++ b/datasette/views/base.py
@@ -20,8 +20,10 @@ from datasette.utils import (
     path_from_row_pks,
     path_with_added_args,
     path_with_format,
+    remove_infinites,
     resolve_table_and_format,
-    to_css_class
+    to_css_class,
+    value_as_boolean,
 )
 
 ureg = pint.UnitRegistry()
@@ -334,6 +336,12 @@ class BaseView(RenderMixin):
                     data["rows"], data["columns"], json_cols,
                 )
 
+            # unless _json_infinity=1 requested, replace infinity with None
+            if "rows" in data and not value_as_boolean(
+                request.args.get("_json_infinity", "0")
+            ):
+                data["rows"] = [remove_infinites(row) for row in data["rows"]]
+
             # Deal with the _shape option
             shape = request.args.get("_shape", "arrays")
             if shape == "arrayfirst":
diff --git a/docs/json_api.rst b/docs/json_api.rst
index b2dc5ff6..4415ceda 100644
--- a/docs/json_api.rst
+++ b/docs/json_api.rst
@@ -148,6 +148,12 @@ querystring arguments:
 
     Compare `this query without the argument <https://fivethirtyeight.datasettes.com/fivethirtyeight.json?sql=select+%27{%22this+is%22%3A+%22a+json+object%22}%27+as+d&_shape=array>`_ to `this query using the argument <https://fivethirtyeight.datasettes.com/fivethirtyeight.json?sql=select+%27{%22this+is%22%3A+%22a+json+object%22}%27+as+d&_shape=array&_json=d>`_
 
+``?_json_infinity=on``
+    If your data contains infinity or -infinity values, Datasette will replace
+    them with None when returning them as JSON. If you pass ``_json_infinity=1``
+    Datasette will instead return them as ``Infinity`` or ``-Infinity`` which is
+    invalid JSON but can be processed by some custom JSON parsers.
+
 ``?_timelimit=MS``
     Sets a custom time limit for the query in ms. You can use this for optimistic
     queries where you would like Datasette to give up if the query takes too
diff --git a/tests/fixtures.py b/tests/fixtures.py
index f2dc7502..6335494a 100644
--- a/tests/fixtures.py
+++ b/tests/fixtures.py
@@ -367,6 +367,15 @@ CREATE TABLE [select] (
 );
 INSERT INTO [select] VALUES ('group', 'having', 'and');
 
+CREATE TABLE infinity (
+    value REAL
+);
+INSERT INTO infinity VALUES
+    (1e999),
+    (-1e999),
+    (1.5)
+;
+
 CREATE TABLE facet_cities (
     id integer primary key,
     name text
diff --git a/tests/test_api.py b/tests/test_api.py
index ac9fb615..8f67a9eb 100644
--- a/tests/test_api.py
+++ b/tests/test_api.py
@@ -18,7 +18,7 @@ def test_homepage(app_client):
     assert response.json.keys() == {'fixtures': 0}.keys()
     d = response.json['fixtures']
     assert d['name'] == 'fixtures'
-    assert d['tables_count'] == 19
+    assert d['tables_count'] == 20
 
 
 def test_database_page(app_client):
@@ -153,6 +153,15 @@ def test_database_page(app_client):
         'label_column': None,
         'fts_table': None,
         'primary_keys': ['pk'],
+    }, {
+        "name": "infinity",
+        "columns": ["value"],
+        "count": 3,
+        "primary_keys": [],
+        "label_column": None,
+        "hidden": False,
+        "fts_table": None,
+        "foreign_keys": {"incoming": [], "outgoing": []}
     }, {
         'columns': ['id', 'content', 'content2'],
         'name': 'primary_key_multiple_columns',
@@ -1281,3 +1290,21 @@ def test_config_force_https_urls():
             "toggle_url"
         ].startswith("https://")
         assert response.json["suggested_facets"][0]["toggle_url"].startswith("https://")
+
+
+def test_infinity_returned_as_null(app_client):
+    response = app_client.get("/fixtures/infinity.json?_shape=array")
+    assert [
+        {"rowid": 1, "value": None},
+        {"rowid": 2, "value": None},
+        {"rowid": 3, "value": 1.5}
+    ] == response.json
+
+
+def test_infinity_returned_as_invalid_json_if_requested(app_client):
+    response = app_client.get("/fixtures/infinity.json?_shape=array&_json_infinity=1")
+    assert [
+        {"rowid": 1, "value": float("inf")},
+        {"rowid": 2, "value": float("-inf")},
+        {"rowid": 3, "value": 1.5}
+    ] == response.json