From 481eb60a1b70f864057e5aa0cec7d6dfba07db49 Mon Sep 17 00:00:00 2001 From: Chris Amico Date: Thu, 3 Feb 2022 15:23:21 -0500 Subject: [PATCH] Make GIS helpers into methods on Database and Table. Remove gis.py. --- docs/python-api.rst | 8 +-- sqlite_utils/cli.py | 2 +- sqlite_utils/db.py | 116 ++++++++++++++++++++++++++++++++++ sqlite_utils/gis.py | 141 ------------------------------------------ sqlite_utils/utils.py | 34 +++++++++- tests/test_cli.py | 3 +- tests/test_gis.py | 61 ++++++++---------- 7 files changed, 179 insertions(+), 186 deletions(-) delete mode 100644 sqlite_utils/gis.py diff --git a/docs/python-api.rst b/docs/python-api.rst index 28ef75d..f2ea954 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -2433,25 +2433,25 @@ Spatialite helpers Initialize Spatialite --------------------- -.. autofunction:: sqlite_utils.gis.init_spatialite +.. automethod:: sqlite_utils.db.Database.init_spatialite .. _python_api_gis_find_spatialite: Finding Spatialite ------------------ -.. autofunction:: sqlite_utils.gis.find_spatialite +.. autofunction:: sqlite_utils.utils.find_spatialite .. _python_api_gis_add_geometry_column: Adding geometry columns ----------------------- -.. autofunction:: sqlite_utils.gis.add_geometry_column +.. automethod:: sqlite_utils.db.Table.add_geometry_column .. _python_api_gis_create_spatial_index: Creating a spatial index ------------------------ -.. autofunction:: sqlite_utils.gis.create_spatial_index +.. automethod:: sqlite_utils.db.Table.create_spatial_index diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 26083b8..f124ffd 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -19,6 +19,7 @@ import tabulate from .utils import ( _compile_code, file_progress, + find_spatialite, sqlite3, decode_base64_values, progressbar, @@ -27,7 +28,6 @@ from .utils import ( TypeTracker, ) -from .gis import find_spatialite CONTEXT_SETTINGS = dict(help_option_names=["-h", "--help"]) diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index e40eb12..7c191b1 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -930,6 +930,43 @@ class Database: sql += " [{}]".format(name) self.execute(sql) + def init_spatialite(self, path: str) -> bool: + """ + The ``init_spatialite`` method will load and initialize the Spatialite extension. + The ``path`` argument should be an absolute path to the compiled extension, which + can be found using ``find_spatialite``. + + Returns true if Spatialite was successfully initialized. + + .. code-block:: python + + from sqlite_utils.db import Database + from sqlite_utils.utils import find_spatialite + + db = Database("mydb.db") + db.init_spatialite(find_spatialite()) + + If you've installed Spatialite somewhere unexpected (for testing an alternate version, for example) + you can pass in an absolute path: + + .. code-block:: python + + from sqlite_utils.db import Database + from sqlite_utils.utils import find_spatialite + + db = Database("mydb.db") + db.init_spatialite("./local/mod_spatialite.dylib") + + """ + self.conn.enable_load_extension(True) + self.conn.load_extension(path) + # Initialize SpatiaLite if not yet initialized + if "spatial_ref_sys" in self.table_names(): + return False + cursor = self.execute("select InitSpatialMetadata(1)") + result = cursor.fetchone() + return result and bool(result[0]) + class Queryable: def exists(self) -> bool: @@ -3011,6 +3048,85 @@ class Table(Queryable): least_common, ) + def add_geometry_column( + self, + column_name: str, + geometry_type: str, + srid: int = 4326, + coord_dimension: str = "XY", + not_null: bool = False, + ) -> bool: + """ + In Spatialite, a geometry column can only be added to an existing table. + To do so, use ``table.add_geometry_column``, passing in a geometry type. + + By default, this will add a nullable column called ``geometry`` using + `SRID 4326 `__. These can be customized using + the ``column_name`` and ``srid`` arguments. + + Returns True if the column was successfully added, False if not. + + .. code-block:: python + + from sqlite_utils.db import Database + from sqlite_utils.utils import find_spatialite + + db = Database("mydb.db") + db.init_spatialite(find_spatialite()) + + # the table must exist before adding a geometry column + table = db["locations"].create({"name": str}) + table.add_geometry_column("geometry", "POINT") + + """ + cursor = self.db.execute( + "SELECT AddGeometryColumn(?, ?, ?, ?, ?, ?);", + [ + self.name, + column_name, + srid, + geometry_type, + coord_dimension, + int(not_null), + ], + ) + + result = cursor.fetchone() + return result and bool(result[0]) + + def create_spatial_index(self, column_name) -> bool: + """ + A spatial index allows for significantly faster bounding box queries. + To create one, use ``create_spatial_index`` with a :ref:`table ` + and the name of an existing geometry column. + + Returns ``True`` if the index was successfully created, ``False`` if not. Calling this + function if an index already exists is a no-op. + + .. code-block:: python + + # assuming Spatialite is loaded, create the table, add the column + table = db["locations"].create({"name": str}) + table.add_geometry_column("geometry", "POINT") + + # now we can index it + table.create_spatial_index("geometry") + + # the spatial index is a virtual table, which we can inspect + print(db["idx_locations_geometry"].schema) + # outputs: + # CREATE VIRTUAL TABLE "idx_locations_geometry" USING rtree(pkid, xmin, xmax, ymin, ymax) + + """ + if f"idx_{self.name}_{column_name}" in self.db.table_names(): + return False + + cursor = self.db.execute( + "select CreateSpatialIndex(?, ?)", [self.name, column_name] + ) + result = cursor.fetchone() + return result and bool(result[0]) + class View(Queryable): def exists(self): diff --git a/sqlite_utils/gis.py b/sqlite_utils/gis.py deleted file mode 100644 index 762c884..0000000 --- a/sqlite_utils/gis.py +++ /dev/null @@ -1,141 +0,0 @@ -import os - -SPATIALITE_PATHS = ( - "/usr/lib/x86_64-linux-gnu/mod_spatialite.so", - "/usr/local/lib/mod_spatialite.dylib", -) - - -def find_spatialite() -> str: - """ - The ``find_spatialite()`` function searches for the `SpatiaLite `__ SQLite extension in some common places. It returns a string path to the location, or ``None`` if SpatiaLite was not found. - - You can use it in code like this: - - .. code-block:: python - - from sqlite_utils import Database - from sqlite_utils.gis import find_spatialite - - db = Database("mydb.db") - spatialite = find_spatialite() - if spatialite: - db.conn.enable_load_extension(True) - db.conn.load_extension(spatialite) - """ - for path in SPATIALITE_PATHS: - if os.path.exists(path): - return path - return None - - -def init_spatialite(db, path: str) -> bool: - """ - The ``init_spatialite`` function will load and initialize the Spatialite extension. - The ``path`` argument should be an absolute path to the compiled extension, which - can be found using ``find_spatialite``. - - Returns true if Spatialite was successfully initialized. - - .. code-block:: python - - from sqlite_utils.gis import find_spatialite, init_spatialite - - db = Database("mydb.db") - init_spatialite(db, find_spatialite()) - - If you've installed Spatialite somewhere unexpected (for testing an alternate version, for example) - you can pass in an absolute path: - - .. code-block:: python - - from sqlite_utils.gis import init_spatialite - - db = Database("mydb.db") - init_spatialite(db, "./local/mod_spatialite.dylib") - - """ - 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 False - cursor = db.execute("select InitSpatialMetadata(1)") - result = cursor.fetchone() - return result and bool(result[0]) - - -def add_geometry_column( - table, - geometry_type: str, - column_name: str = "geometry", - srid: int = 4326, - coord_dimension: str = "XY", - not_null: bool = False, -) -> bool: - """ - In Spatialite, a geometry column can only be added to an existing table. - To do so, use ``add_geometry_column``, passing in a :ref:`table ` - and geometry type. - - By default, this will add a nullable column called ``geometry`` using - `SRID 4326 `__. These can be customized using - the ``column_name`` and ``srid`` arguments. - - Returns True if the column was successfully added, False if not. - - .. code-block:: python - - from sqlite_utils.gis import find_spatialite, init_spatialite, add_geometry_column - - db = Database("mydb.db") - init_spatialite(db, find_spatialite()) - - # the table must exist before adding a geometry column - db["locations"].create({"name": str}) - add_geometry_column(db["locations"], "POINT") - - """ - cursor = table.db.execute( - "SELECT AddGeometryColumn(?, ?, ?, ?, ?, ?);", - [table.name, column_name, srid, geometry_type, coord_dimension, int(not_null)], - ) - - result = cursor.fetchone() - return result and bool(result[0]) - - -def create_spatial_index(table, column_name: str = "geometry") -> bool: - """ - A spatial index allows for significantly faster bounding box queries. - To create on, use ``create_spatial_index`` with a :ref:`table ` - and the name of an existing geometry column. - - Returns True if the index was successfully created, False if not. Calling this - function if an index already exists is a no-op. - - .. code-block:: python - - from sqlite_utils.gis import add_geometry_column, create_spatial_index - - # assuming Spatialite is loaded, create the table, add the column - db["locations"].create({"name": str}) - add_geometry_column(db["locations"], "POINT", "geometry") - - # now we can index it - create_spatial_index(db["locations"], "geometry") - - # the spatial index is a virtual table, which we can inspect - print(db["idx_locations_geometry"].schema) - # outputs: - # CREATE VIRTUAL TABLE "idx_locations_geometry" USING rtree(pkid, xmin, xmax, ymin, ymax) - - """ - if f"idx_{table.name}_{column_name}" in table.db.table_names(): - return False - - cursor = table.db.execute( - "select CreateSpatialIndex(?, ?)", [table.name, column_name] - ) - result = cursor.fetchone() - return result and bool(result[0]) diff --git a/sqlite_utils/utils.py b/sqlite_utils/utils.py index 6fdc04f..c2a7c91 100644 --- a/sqlite_utils/utils.py +++ b/sqlite_utils/utils.py @@ -22,8 +22,38 @@ except ImportError: OperationalError = sqlite3.OperationalError -# backwards compatibility -from .gis import find_spatialite + +SPATIALITE_PATHS = ( + "/usr/lib/x86_64-linux-gnu/mod_spatialite.so", + "/usr/local/lib/mod_spatialite.dylib", +) + + +def find_spatialite() -> str: + """ + The ``find_spatialite()`` function searches for the `SpatiaLite `__ SQLite extension in some common places. It returns a string path to the location, or ``None`` if SpatiaLite was not found. + + You can use it in code like this: + + .. code-block:: python + + from sqlite_utils import Database + from sqlite_utils.utils import find_spatialite + + db = Database("mydb.db") + spatialite = find_spatialite() + if spatialite: + db.conn.enable_load_extension(True) + db.conn.load_extension(spatialite) + + # or use with db.init_spatialite like this + db.init_spatialite(find_spatialite()) + + """ + for path in SPATIALITE_PATHS: + if os.path.exists(path): + return path + return None def suggest_column_types(records): diff --git a/tests/test_cli.py b/tests/test_cli.py index 3dc097e..60c95b0 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -7,8 +7,7 @@ from unittest import mock import json import os import pytest -from sqlite_utils.utils import sqlite3 -from sqlite_utils.gis import find_spatialite +from sqlite_utils.utils import sqlite3, find_spatialite import textwrap from .utils import collapse_whitespace diff --git a/tests/test_gis.py b/tests/test_gis.py index afb9a39..90804a3 100644 --- a/tests/test_gis.py +++ b/tests/test_gis.py @@ -1,53 +1,46 @@ import pytest -from sqlite_utils import gis +from sqlite_utils.utils import find_spatialite 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 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() + spatialite = 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 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) + spatialite = find_spatialite() + db.init_spatialite(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 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) + spatialite = find_spatialite() + db.init_spatialite(spatialite) # create a table first - db.create_table("locations", {"id": str, "properties": str}) - gis.add_geometry_column( - db["locations"], - geometry_type="Point", + table = db.create_table("locations", {"id": str, "properties": str}) + table.add_geometry_column( column_name="geometry", + geometry_type="Point", srid=4326, coord_dimension=2, ) @@ -62,48 +55,44 @@ def test_add_geometry_column(): } -@pytest.mark.skipif( - not gis.find_spatialite(), reason="Could not find SpatiaLite extension" -) +@pytest.mark.skipif(not 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() - assert gis.init_spatialite(db, spatialite) + spatialite = find_spatialite() + assert db.init_spatialite(spatialite) # create a table, add a geometry column with default values - db.create_table("locations", {"id": str, "properties": str}) - assert gis.add_geometry_column(db["locations"], "Point", "geometry") + table = db.create_table("locations", {"id": str, "properties": str}) + assert table.add_geometry_column("geometry", "Point") # index it - assert gis.create_spatial_index(db["locations"], "geometry") + assert table.create_spatial_index("geometry") assert "idx_locations_geometry" in db.table_names() -@pytest.mark.skipif( - not gis.find_spatialite(), reason="Could not find SpatiaLite extension" -) +@pytest.mark.skipif(not 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_double_create_spatial_index(): db = Database(memory=True) - spatialite = gis.find_spatialite() - gis.init_spatialite(db, spatialite) + spatialite = find_spatialite() + db.init_spatialite(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") + table = db.create_table("locations", {"id": str, "properties": str}) + table.add_geometry_column("geometry", "Point") # index it, return True - assert gis.create_spatial_index(db["locations"], "geometry") + assert table.create_spatial_index("geometry") assert "idx_locations_geometry" in db.table_names() # call it again, return False - assert not gis.create_spatial_index(db["locations"], "geometry") + assert not table.create_spatial_index("geometry")