diff --git a/geojson_pydantic/features.py b/geojson_pydantic/features.py index 071b340..90620b7 100644 --- a/geojson_pydantic/features.py +++ b/geojson_pydantic/features.py @@ -1,19 +1,23 @@ """pydantic models for GeoJSON Feature objects.""" -from typing import Any, Dict, List, Optional +from typing import Dict, Generic, List, Optional, TypeVar -from pydantic import BaseModel, Field, validator +from pydantic import Field, validator +from pydantic.generics import GenericModel from .geometries import Geometry from .utils import BBox +Props = TypeVar("Props", bound=Dict) +Geom = TypeVar("Geom", bound=Geometry) -class Feature(BaseModel): + +class Feature(GenericModel, Generic[Geom, Props]): """Feature Model""" type: str = Field("Feature", const=True) - geometry: Geometry - properties: Optional[Dict[Any, Any]] + geometry: Geom + properties: Optional[Props] id: Optional[str] bbox: Optional[BBox] @@ -30,11 +34,11 @@ def set_geometry(cls, v): return v -class FeatureCollection(BaseModel): +class FeatureCollection(GenericModel, Generic[Geom, Props]): """FeatureCollection Model""" type: str = Field("FeatureCollection", const=True) - features: List[Feature] + features: List[Feature[Geom, Props]] bbox: Optional[BBox] def __iter__(self): diff --git a/test/test_features.py b/test/test_features.py index a01f623..f39ea7b 100644 --- a/test/test_features.py +++ b/test/test_features.py @@ -1,14 +1,94 @@ +from random import randint +from typing import Dict +from uuid import uuid4 + import pytest +from pydantic import BaseModel, ValidationError from geojson_pydantic.features import Feature, FeatureCollection -from geojson_pydantic.geometries import Point +from geojson_pydantic.geometries import Geometry, MultiPolygon, Polygon + + +class GenericProperties(BaseModel): + id: str + description: str + size: int + + +properties = { + "id": str(uuid4()), + "description": str(uuid4()), + "size": randint(0, 1000), +} +polygon = { + "type": "Polygon", + "coordinates": [ + [ + [13.38272, 52.46385], + [13.42786, 52.46385], + [13.42786, 52.48445], + [13.38272, 52.48445], + [13.38272, 52.46385], + ] + ], +} -@pytest.mark.parametrize("coordinates", [(1, 2)]) -def test_geometry_collection_iteration(coordinates): +test_feature = { + "type": "Feature", + "geometry": polygon, + "properties": properties, +} + + +def test_geometry_collection_iteration(): """test if feature collection is iterable""" - print(coordinates, type(coordinates)) - polygon = Point(coordinates=coordinates) - feature = Feature(geometry=polygon) - gc = FeatureCollection(features=[feature, feature]) + gc = FeatureCollection(features=[test_feature, test_feature]) iter(gc) + + +def test_generic_properties_is_dict(): + feature = Feature(**test_feature) + assert feature.properties["id"] == test_feature["properties"]["id"] + assert type(feature.properties) == dict + assert not hasattr(feature.properties, "id") + + +def test_generic_properties_is_object(): + feature = Feature[Geometry, GenericProperties](**test_feature) + assert feature.properties.id == test_feature["properties"]["id"] + assert type(feature.properties) == GenericProperties + assert hasattr(feature.properties, "id") + + +def test_generic_geometry(): + feature = Feature[Polygon, GenericProperties](**test_feature) + assert feature.properties.id == test_feature["properties"]["id"] + assert type(feature.geometry) == Polygon + assert type(feature.properties) == GenericProperties + assert hasattr(feature.properties, "id") + + feature = Feature[Polygon, Dict](**test_feature) + assert type(feature.geometry) == Polygon + assert feature.properties["id"] == test_feature["properties"]["id"] + assert type(feature.properties) == dict + assert not hasattr(feature.properties, "id") + + with pytest.raises(ValidationError): + Feature[MultiPolygon, Dict](**({"type": "Feature", "geometry": polygon})) + + +def test_generic_properties_should_raise_for_string(): + with pytest.raises(ValidationError): + Feature( + **({"type": "Feature", "geometry": polygon, "properties": "should raise"}) + ) + + +def test_feature_collection_generic(): + fc = FeatureCollection[Polygon, GenericProperties]( + features=[test_feature, test_feature] + ) + assert len(fc) == 2 + assert type(fc[0].properties) == GenericProperties + assert type(fc[0].geometry) == Polygon