add count command

Author: smnorris

.pre-commit-config.yaml

@@ -4,6 +4,7 @@ repos:
   hooks:
     - id: ruff
       name: lint with ruff
+      args: [--fix]
     - id: ruff
       name: sort imports with ruff
       args: [--select, I, --fix]

src/bcdata/cli.py

@@ -6,12 +6,6 @@
 
 import click
 from cligj import compact_opt, indent_opt, quiet_opt, verbose_opt
-from shapely.geometry.linestring import LineString
-from shapely.geometry.multilinestring import MultiLineString
-from shapely.geometry.multipoint import MultiPoint
-from shapely.geometry.multipolygon import MultiPolygon
-from shapely.geometry.point import Point
-from shapely.geometry.polygon import Polygon
 
 import bcdata
 from bcdata.database import Database
@@ -27,6 +21,7 @@ def configure_logging(verbosity):
 def complete_dataset_names(ctx, param, incomplete):
     return [k for k in bcdata.list_tables() if k.startswith(incomplete)]
 
+
 # bounds handling direct from rasterio
 # https://github.com/mapbox/rasterio/blob/master/rasterio/rio/options.py
 # https://github.com/mapbox/rasterio/blob/master/rasterio/rio/clip.py
@@ -202,6 +197,13 @@ def dem(
     help="A valid CQL or ECQL query",
 )
 @click.option("--out_file", "-o", help="Output file")
+@click.option(
+    "--count",
+    "-c",
+    default=None,
+    type=int,
+    help="Number of features to request and dump",
+)
 @bounds_opt
 @click.option(
     "--bounds-crs",
@@ -219,7 +221,7 @@ def dem(
 @lowercase_opt
 @verbose_opt
 @quiet_opt
-def dump(dataset, query, out_file, bounds, bounds_crs, no_clean, lowercase, verbose, quiet):
+def dump(dataset, query, out_file, count, bounds, bounds_crs, no_clean, lowercase, verbose, quiet):
     """Write DataBC features to stdout as GeoJSON feature collection.
 
     \b
@@ -240,7 +242,13 @@ def dump(dataset, query, out_file, bounds, bounds_crs, no_clean, lowercase, verb
     else:
         clean = True
     data = bcdata.get_data(
-        table, query=query, bounds=bounds, bounds_crs=bounds_crs, lowercase=lowercase, clean=clean
+        table,
+        query=query,
+        count=count,
+        bounds=bounds,
+        bounds_crs=bounds_crs,
+        lowercase=lowercase,
+        clean=clean,
     )
     if out_file:
         with open(out_file, "w") as sink:

tests/test_cli.py

@@ -1,3 +1,4 @@
+import json
 import os
 
 from click.testing import CliRunner
@@ -73,6 +74,14 @@ def test_cat_bounds_ll():
     assert len(result.output.split("\n")) == 4
 
 
+def test_dump():
+    runner = CliRunner()
+    result = runner.invoke(cli, ["dump", AIRPORTS_TABLE, "--count", 1])
+    assert result.exit_code == 0
+    assert json.loads(result.output)["type"] == "FeatureCollection"
+    assert len(json.loads(result.output)["features"]) == 1
+
+
 def test_bc2pg():
     runner = CliRunner()
     result = runner.invoke(cli, ["bc2pg", AIRPORTS_TABLE, "--db_url", DB_URL])

tests/test_wfs.py

@@ -1,5 +1,4 @@
 import pytest
-import json
 import requests
 import requests_mock
 import stamina
@@ -94,9 +93,7 @@ def test_get_data_lowercase():
 
 def test_get_data_crs():
     data = bcdata.get_data(AIRPORTS_TABLE, crs="EPSG:3005")
-    assert (
-        data["crs"]["properties"]["name"] == 'urn:ogc:def:crs:EPSG::3005'
-    )
+    assert data["crs"]["properties"]["name"] == "urn:ogc:def:crs:EPSG::3005"
 
 
 def test_get_features():
@@ -159,8 +156,6 @@ def test_clean():
 
 def test_no_clean():
     data = bcdata.get_data(
-        AIRPORTS_TABLE,
-        query="AIRPORT_NAME='Terrace (Northwest Regional) Airport'",
-        clean=False
+        AIRPORTS_TABLE, query="AIRPORT_NAME='Terrace (Northwest Regional) Airport'", clean=False
     )
     assert "SE_ANNO_CAD_DATA" in data["features"][0]["properties"].keys()