add helper methods to field extension (#69)

* add helper methods to field extension

* add test

Author: geospatial-jeff

stac_pydantic/api/extensions/fields.py

@@ -1,4 +1,4 @@
-from typing import Optional, Set
+from typing import Dict, Optional, Set
 
 from pydantic import BaseModel
 
@@ -10,3 +10,41 @@ class FieldsExtension(BaseModel):
 
     includes: Optional[Set[str]]
     excludes: Optional[Set[str]]
+
+    def _get_field_dict(self, fields: Set[str]) -> Dict:
+        """Internal method to create a dictionary for advanced include or exclude of pydantic fields on model export
+
+        Ref: https://pydantic-docs.helpmanual.io/usage/exporting_models/#advanced-include-and-exclude
+        """
+        field_dict = {}
+        for field in fields:
+            if "." in field:
+                parent, key = field.split(".")
+                if parent not in field_dict:
+                    field_dict[parent] = {key}
+                else:
+                    field_dict[parent].add(key)
+            else:
+                field_dict[field] = ...  # type:ignore
+        return field_dict
+
+    @property
+    def filter(self) -> Dict:
+        """
+        Create dictionary of fields to include/exclude on model export based on the included and excluded fields passed
+        to the API.  The output of this property may be passed to pydantic's serialization methods to include or exclude
+        certain keys.
+
+        Ref: https://pydantic-docs.helpmanual.io/usage/exporting_models/#advanced-include-and-exclude
+        """
+        include = set()
+        # If only include is specified, add fields to the set
+        if self.includes and not self.excludes:
+            include = include.union(self.includes)
+        # If both include + exclude specified, find the difference between sets
+        elif self.includes and self.excludes:
+            include = include.union(self.includes) - self.excludes
+        return {
+            "include": self._get_field_dict(include),
+            "exclude": self._get_field_dict(self.excludes),
+        }

tests/test_api_extensions.py

@@ -0,0 +1,31 @@
+from datetime import datetime
+
+import pytest
+from shapely.geometry import Polygon
+
+from stac_pydantic import Item
+from stac_pydantic.api.extensions.fields import FieldsExtension
+
+
+def test_fields_filter():
+    fields = FieldsExtension(
+        includes={"id", "geometry", "properties.foo"}, excludes={"properties.bar"}
+    )
+
+    item = Item(
+        id="test-fields-filter",
+        geometry=Polygon.from_bounds(0, 0, 0, 0),
+        properties={"datetime": datetime.utcnow(), "foo": "foo", "bar": "bar"},
+        assets={},
+        links=[],
+        bbox=[0, 0, 0, 0],
+    )
+
+    d = item.to_dict(**fields.filter)
+    assert d.pop("id") == item.id
+    assert d.pop("geometry") == item.geometry
+    props = d.pop("properties")
+    assert props["foo"] == "foo"
+
+    assert not props.get("bar")
+    assert not d