Add new gis.py module with spatialite helper methods

docs/python-api.rst

@@ -2352,7 +2352,7 @@ You can use it in code like this:
 .. code-block:: python
 
     from sqlite_utils import Database
-    from sqlite_utils.utils import find_spatialite
+    from sqlite_utils.gis import find_spatialite
 
     db = Database("mydb.db")
     spatialite = find_spatialite()

sqlite_utils/cli.py

@@ -19,7 +19,6 @@ import tabulate
 from .utils import (
     _compile_code,
     file_progress,
-    find_spatialite,
     sqlite3,
     decode_base64_values,
     progressbar,
@@ -28,6 +27,8 @@ from .utils import (
     TypeTracker,
 )
 
+from .gis import find_spatialite
+
 CONTEXT_SETTINGS = dict(help_option_names=["-h", "--help"])
 
 VALID_COLUMN_TYPES = ("INTEGER", "TEXT", "FLOAT", "BLOB")

sqlite_utils/gis.py

@@ -0,0 +1,44 @@
+import os
+from .db import Database, Table
+
+SPATIALITE_PATHS = (
+    "/usr/lib/x86_64-linux-gnu/mod_spatialite.so",
+    "/usr/local/lib/mod_spatialite.dylib",
+)
+
+
+def find_spatialite() -> str:
+    for path in SPATIALITE_PATHS:
+        if os.path.exists(path):
+            return path
+    return None
+
+
+def init_spatialite(db: Database, path: str) -> None:
+    "Load spatialite extension for a database"
+    db.conn.enable_load_extension(True)
+    db.conn.load_extension(path)
+    # Initialize SpatiaLite if not yet initialized
+    if "spatial_ref_sys" in db.table_names():
+        return
+    db.conn.execute("select InitSpatialMetadata(1)")
+
+
+def add_geometry_column(
+    table: Table,
+    geometry_type: str,
+    column_name: str = "geometry",
+    srid: int = 4326,
+    coord_dimension: str = "XY",
+    not_null: bool = False,
+) -> None:
+    "Add a geometry column to a table"
+    table.db.conn.execute(
+        "SELECT AddGeometryColumn(?, ?, ?, ?, ?, ?);",
+        [table.name, column_name, srid, geometry_type, coord_dimension, int(not_null)],
+    )
+
+
+def create_spatial_index(table: Table, column_name: str = "geometry") -> None:
+    "Create a spatial index for a table and column"
+    table.db.conn.execute("select CreateSpatialIndex(?, ?)", [table.name, column_name])

sqlite_utils/utils.py

@@ -22,11 +22,6 @@ except ImportError:
 
     OperationalError = sqlite3.OperationalError
 
-SPATIALITE_PATHS = (
-    "/usr/lib/x86_64-linux-gnu/mod_spatialite.so",
-    "/usr/local/lib/mod_spatialite.dylib",
-)
-
 
 def suggest_column_types(records):
     all_column_types = {}
@@ -96,13 +91,6 @@ def decode_base64_values(doc):
     return dict(doc, **{k: base64.b64decode(doc[k]["encoded"]) for k in to_fix})
 
 
-def find_spatialite():
-    for path in SPATIALITE_PATHS:
-        if os.path.exists(path):
-            return path
-    return None
-
-
 class UpdateWrapper:
     def __init__(self, wrapped, update):
         self._wrapped = wrapped

tests/test_cli.py

@@ -7,7 +7,8 @@ from unittest import mock
 import json
 import os
 import pytest
-from sqlite_utils.utils import sqlite3, find_spatialite
+from sqlite_utils.utils import sqlite3
+from sqlite_utils.gis import find_spatialite
 import textwrap
 
 from .utils import collapse_whitespace

tests/test_gis.py

@@ -0,0 +1,167 @@
+import pytest
+from sqlite_utils import gis
+from sqlite_utils.db import Database
+from sqlite_utils.utils import sqlite3
+
+
+@pytest.mark.skipif(
+    not gis.find_spatialite(), reason="Could not find SpatiaLite extension"
+)
+@pytest.mark.skipif(
+    not hasattr(sqlite3.Connection, "enable_load_extension"),
+    reason="sqlite3.Connection missing enable_load_extension",
+)
+def test_find_spatialite():
+    spatialite = gis.find_spatialite()
+    assert spatialite is None or isinstance(spatialite, str)
+
+
+@pytest.mark.skipif(
+    not gis.find_spatialite(), reason="Could not find SpatiaLite extension"
+)
+@pytest.mark.skipif(
+    not hasattr(sqlite3.Connection, "enable_load_extension"),
+    reason="sqlite3.Connection missing enable_load_extension",
+)
+def test_init_spatialite():
+    db = Database(memory=True)
+    spatialite = gis.find_spatialite()
+    gis.init_spatialite(db, spatialite)
+    assert "spatial_ref_sys" in db.table_names()
+
+
+@pytest.mark.skipif(
+    not gis.find_spatialite(), reason="Could not find SpatiaLite extension"
+)
+@pytest.mark.skipif(
+    not hasattr(sqlite3.Connection, "enable_load_extension"),
+    reason="sqlite3.Connection missing enable_load_extension",
+)
+def test_add_geometry_column():
+    db = Database(memory=True)
+    spatialite = gis.find_spatialite()
+    gis.init_spatialite(db, spatialite)
+
+    # create a table first
+    db.create_table("locations", {"id": str, "properties": str})
+    gis.add_geometry_column(
+        db["locations"],
+        geometry_type="Point",
+        column_name="geometry",
+        srid=4326,
+        coord_dimension=2,
+    )
+
+    assert db["geometry_columns"].get(["locations", "geometry"]) == {
+        "f_table_name": "locations",
+        "f_geometry_column": "geometry",
+        "geometry_type": 1,  # point
+        "coord_dimension": 2,
+        "srid": 4326,
+        "spatial_index_enabled": 0,
+    }
+
+
+@pytest.mark.skipif(
+    not gis.find_spatialite(), reason="Could not find SpatiaLite extension"
+)
+@pytest.mark.skipif(
+    not hasattr(sqlite3.Connection, "enable_load_extension"),
+    reason="sqlite3.Connection missing enable_load_extension",
+)
+def test_create_spatial_index():
+    db = Database(memory=True)
+    spatialite = gis.find_spatialite()
+    gis.init_spatialite(db, spatialite)
+
+    # create a table, add a geometry column with default values
+    db.create_table("locations", {"id": str, "properties": str})
+    gis.add_geometry_column(db["locations"], "Point", "geometry")
+
+    # index it
+    gis.create_spatial_index(db["locations"], "geometry")
+
+    assert "idx_locations_geometry" in db.table_names()
+
+
+# extract of In-N-Out locations via alltheplaces.xyz
+locations = {
+    "type": "FeatureCollection",
+    "features": [
+        {
+            "type": "Feature",
+            "id": "APyEi_bhUMkAPx6J_I1JrqB5eBo=",
+            "properties": {
+                "ref": "322",
+                "@spider": "innout",
+                "addr:full": "2835 W. University Dr.",
+                "addr:city": "Denton",
+                "addr:state": "TX",
+                "addr:postcode": "76201",
+                "name": "Denton",
+                "website": "http://locations.in-n-out.com/322",
+            },
+            "geometry": {"type": "Point", "coordinates": [-97.17114, 33.2294]},
+        },
+        {
+            "type": "Feature",
+            "id": "_fK8gYjGcdlGA09B3a4a8RyUZD4=",
+            "properties": {
+                "ref": "255",
+                "@spider": "innout",
+                "addr:full": "190 E. Stacy Rd.",
+                "addr:city": "Allen",
+                "addr:state": "TX",
+                "addr:postcode": "75002",
+                "name": "Allen",
+                "website": "http://locations.in-n-out.com/255",
+            },
+            "geometry": {"type": "Point", "coordinates": [-96.65328, 33.12914]},
+        },
+        {
+            "type": "Feature",
+            "id": "aO9MWd_7HNlDTGdFH-Vl7vscZdk=",
+            "properties": {
+                "ref": "256",
+                "@spider": "innout",
+                "addr:full": "2800 Preston Rd.",
+                "addr:city": "Frisco",
+                "addr:state": "TX",
+                "addr:postcode": "75034",
+                "name": "Frisco",
+                "website": "http://locations.in-n-out.com/256",
+            },
+            "geometry": {"type": "Point", "coordinates": [-96.80456, 33.1018]},
+        },
+        {
+            "type": "Feature",
+            "id": "yHaQ--D7A6FXdx-vrLKdVxctV5g=",
+            "properties": {
+                "ref": "299",
+                "@spider": "innout",
+                "addr:full": "5298 State Highway 121",
+                "addr:city": "The Colony",
+                "addr:state": "TX",
+                "addr:postcode": "75056",
+                "name": "The Colony",
+                "website": "http://locations.in-n-out.com/299",
+            },
+            "geometry": {"type": "Point", "coordinates": [-96.87604, 33.07076]},
+        },
+        {
+            "type": "Feature",
+            "id": "sGpWMAPrvimcDCBoE9xO7-tkLmw=",
+            "properties": {
+                "ref": "329",
+                "@spider": "innout",
+                "addr:full": "3500 Highway 114",
+                "addr:city": "Fort Worth",
+                "addr:state": "TX",
+                "addr:postcode": "76177",
+                "name": "Fort Worth",
+                "website": "http://locations.in-n-out.com/329",
+            },
+            "geometry": {"type": "Point", "coordinates": [-97.28667, 33.02593]},
+        },
+    ],
+}

tests/test_utils.py

@@ -22,11 +22,6 @@ def test_decode_base64_values(input, expected, should_be_is):
         assert actual == expected
 
 
-def test_find_spatialite():
-    spatialite = utils.find_spatialite()
-    assert spatialite is None or isinstance(spatialite, str)
-
-
 @pytest.mark.parametrize(
     "size,expected",
     (