Format python

This commit is contained in:
Andras Schmelczer 2026-01-31 13:07:09 +00:00
parent 85f5770e09
commit 4c258018c3
17 changed files with 348 additions and 248 deletions

View file

@ -7,28 +7,32 @@ from pipeline.utils.poi_counts import POI_GROUPS, count_pois_within_radius
@pytest.fixture
def pois():
"""POIs clustered around two locations: central London and 10km away."""
return pl.DataFrame({
"lat": [51.5074, 51.5075, 51.5080, 51.5076, 51.5073, 51.60],
"lng": [-0.1278, -0.1280, -0.1275, -0.1279, -0.1277, -0.20],
"category": [
"Restaurant",
"Fast Food",
"Supermarket",
"Park",
"Station",
"Restaurant", # too far from any property
],
})
return pl.DataFrame(
{
"lat": [51.5074, 51.5075, 51.5080, 51.5076, 51.5073, 51.60],
"lng": [-0.1278, -0.1280, -0.1275, -0.1279, -0.1277, -0.20],
"category": [
"Restaurant",
"Fast Food",
"Supermarket",
"Park",
"Station",
"Restaurant", # too far from any property
],
}
)
@pytest.fixture
def properties():
"""Two properties at the same postcode near central London, one at a distant postcode."""
return pl.DataFrame({
"postcode": ["EC1A 1BB", "EC1A 1BB", "ZZ99 9ZZ"],
"lat": [51.5074, 51.5074, 55.0],
"lon": [-0.1278, -0.1278, -3.0],
})
return pl.DataFrame(
{
"postcode": ["EC1A 1BB", "EC1A 1BB", "ZZ99 9ZZ"],
"lat": [51.5074, 51.5074, 55.0],
"lon": [-0.1278, -0.1278, -3.0],
}
)
def test_counts_pois_within_radius(properties, pois):
@ -41,9 +45,9 @@ def test_counts_pois_within_radius(properties, pois):
assert len(series) == 3, f"{col} has {len(series)} rows, expected 3"
# First two rows share a postcode near the central London cluster
assert result["restaurants_2km"][0] == 2 # Restaurant + Fast Food
assert result["groceries_2km"][0] == 1 # Supermarket
assert result["parks_2km"][0] == 1 # Park
assert result["restaurants_2km"][0] == 2 # Restaurant + Fast Food
assert result["groceries_2km"][0] == 1 # Supermarket
assert result["parks_2km"][0] == 1 # Park
assert result["public_transport_2km"][0] == 1 # Station
# Second row is the same postcode, so same counts
@ -55,11 +59,13 @@ def test_counts_pois_within_radius(properties, pois):
def test_no_pois_returns_zeros(properties):
empty_pois = pl.DataFrame({
"lat": pl.Series([], dtype=pl.Float64),
"lng": pl.Series([], dtype=pl.Float64),
"category": pl.Series([], dtype=pl.String),
})
empty_pois = pl.DataFrame(
{
"lat": pl.Series([], dtype=pl.Float64),
"lng": pl.Series([], dtype=pl.Float64),
"category": pl.Series([], dtype=pl.String),
}
)
result = count_pois_within_radius(properties, empty_pois, radius_km=2.0)
for group in POI_GROUPS:
@ -70,11 +76,13 @@ def test_no_pois_returns_zeros(properties):
def test_custom_radius(pois):
"""A tiny radius should exclude POIs that are even slightly away."""
properties = pl.DataFrame({
"postcode": ["EC1A 1BB"],
"lat": [51.5074],
"lon": [-0.1278],
})
properties = pl.DataFrame(
{
"postcode": ["EC1A 1BB"],
"lat": [51.5074],
"lon": [-0.1278],
}
)
# 0.01 km = 10m — only the POI at the exact same location should match
result = count_pois_within_radius(properties, pois, radius_km=0.01)