From b579a9c8ae5cc533d3807b5386a06f199c2af1d1 Mon Sep 17 00:00:00 2001 From: Chris Amico Date: Thu, 13 Jan 2022 22:52:08 -0500 Subject: [PATCH] Add new gis.py module with spatialite helper methods --- docs/python-api.rst | 2 +- sqlite_utils/cli.py | 3 +- sqlite_utils/gis.py | 44 +++++++++++ sqlite_utils/utils.py | 12 --- tests/test_cli.py | 3 +- tests/test_gis.py | 167 ++++++++++++++++++++++++++++++++++++++++++ tests/test_utils.py | 5 -- 7 files changed, 216 insertions(+), 20 deletions(-) create mode 100644 sqlite_utils/gis.py create mode 100644 tests/test_gis.py diff --git a/docs/python-api.rst b/docs/python-api.rst index 6b0d542..8ede326 100644 --- a/docs/python-api.rst +++ b/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() diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 9f1331d..26083b8 100644 --- a/sqlite_utils/cli.py +++ b/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") diff --git a/sqlite_utils/gis.py b/sqlite_utils/gis.py new file mode 100644 index 0000000..632b096 --- /dev/null +++ b/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]) diff --git a/sqlite_utils/utils.py b/sqlite_utils/utils.py index 0777f30..df5b439 100644 --- a/sqlite_utils/utils.py +++ b/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 diff --git a/tests/test_cli.py b/tests/test_cli.py index 60c95b0..3dc097e 100644 --- a/tests/test_cli.py +++ b/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 diff --git a/tests/test_gis.py b/tests/test_gis.py new file mode 100644 index 0000000..57ef6ba --- /dev/null +++ b/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]}, + }, + ], +} diff --git a/tests/test_utils.py b/tests/test_utils.py index 8630e28..66ce413 100644 --- a/tests/test_utils.py +++ b/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", (