Merge pull request #111 from predicthq/BU-24-Update-SDK-for-Beam-Saved-Location-support

Bu 24 update sdk for beam saved location support

Author: PaulDDD

predicthq/endpoints/v1/beam/endpoint.py

@@ -8,6 +8,8 @@
     CreateAnalysisGroupResponse,
     AnalysisGroup,
     AnalysisGroupResultSet,
+    GeoJson,
+    Place
 )
 from predicthq.endpoints.decorators import accepts, returns
 from typing import overload, List, Optional, TextIO, Union
@@ -29,6 +31,7 @@ def create(
             location__unit: Optional[str] = None,
             location__google_place_id: Optional[str] = None,
             location__geoscope_paths: Optional[List[str]] = None,
+            location__saved_location_id: Optional[str] = None,
             rank__type: Optional[str] = None,
             rank__levels__phq: Optional[dict] = None,
             rank__levels__local: Optional[dict] = None,
@@ -70,6 +73,7 @@ def search(
             limit: Optional[int] = None,
             external_id: Optional[List[str]] = None,
             label: Optional[List[str]] = None,
+            location__saved_location_id: Optional[List[str]] = None,
             **params,
         ): ...
         @accepts()
@@ -102,6 +106,7 @@ def update(
             location__unit: Optional[str] = None,
             location__google_place_id: Optional[str] = None,
             location__geoscope_paths: Optional[List[str]] = None,
+            location__saved_location_id: Optional[str] = None,
             rank__type: Optional[str] = None,
             rank__levels__phq: Optional[dict] = None,
             rank__levels__local: Optional[dict] = None,

predicthq/endpoints/v1/beam/schemas.py

@@ -3,6 +3,28 @@
 from predicthq.endpoints.schemas import ArgKwargResultSet
 from typing import Optional, List
 
+# Python < 3.11 does not have StrEnum in the enum module
+import sys
+if sys.version_info < (3, 11):
+    import enum
+
+    class StrEnum(str, enum.Enum):
+        pass
+else:
+    from enum import StrEnum
+
+# Python < 3.9 does not have Annotated
+if sys.version_info < (3, 9):
+    from typing_extensions import Annotated
+else:
+    from typing import Annotated
+
+# Python < 3.8 does not have Literal
+if sys.version_info < (3, 8):
+    from typing_extensions import Literal
+else:
+    from typing import Literal
+
 
 class BeamPaginationResultSet(ArgKwargResultSet):
     def has_next(self):
@@ -113,6 +135,58 @@ class DemandType(DemandTypeGroup):
     currency_code: str
 
 
+class RadiusUnit(StrEnum):
+    m = "m"
+    km = "km"
+    mi = "mi"
+    ft = "ft"
+
+
+class GeoJsonGeometryType(StrEnum):
+    POINT = "Point"
+    POLYGON = "Polygon"
+    MULTI_POLYGON = "MultiPolygon"
+    LINE_STRING = "LineString"
+    MULTI_LINE_STRING = "MultiLineString"
+
+
+class GeoJsonProperties(BaseModel):
+    radius: Annotated[float, Field(gt=0)]
+    radius_unit: RadiusUnit
+
+
+class GeoJsonGeometry(BaseModel):
+    type: GeoJsonGeometryType
+    coordinates: Annotated[list, Field(min_length=1)]
+
+
+class GeoJson(BaseModel):
+    type: Literal["Feature"]
+    properties: Optional[GeoJsonProperties] = None
+    geometry: GeoJsonGeometry
+
+
+class Place(BaseModel):
+    place_id: int
+    type: str
+    name: str
+    county: Optional[str] = None
+    region: Optional[str] = None
+    country: Optional[str] = None
+    geojson: GeoJson
+
+
+class SavedLocation(BaseModel):
+    name: Optional[str] = None
+    formatted_address: Optional[str] = None
+    geojson: Optional[GeoJson] = None
+    h3: Optional[List[str]] = None
+    place_ids: Optional[List[int]] = None
+    place_hierarchies: Optional[List[str]] = None
+    places: Optional[List[Place]] = None
+    location_id: str
+
+
 class Analysis(BaseModel):
     model_config: ConfigDict = ConfigDict(extra="allow")
 
@@ -134,6 +208,7 @@ class Analysis(BaseModel):
     processed_dt: Optional[datetime] = None
     external_id: Optional[str] = None
     label: Optional[List[str]] = None
+    saved_location: Optional[SavedLocation] = None
 
 
 class AnalysisResultSet(BeamPaginationResultSet):

tests/endpoints/v1/test_beam.py

@@ -50,6 +50,7 @@ def test_search_params_underscores(self, client):
             limit=10,
             external_id="external_id",
             label=["label1", "label2"],
+            location__saved_location_id="saved_location_id",
         )
 
         client.request.assert_called_once_with(
@@ -72,6 +73,7 @@ def test_search_params_underscores(self, client):
                 "limit": 10,
                 "external_id": "external_id",
                 "label": "label1,label2",
+                "location.saved_location_id": "saved_location_id",
             },
             verify=True,
         )
@@ -98,6 +100,7 @@ def test_search_params_dicts(self, client):
             limit=10,
             external_id="external_id",
             label=["label1", "label2"],
+            location__saved_location_id="saved_location_id",
         )
 
         client.request.assert_called_once_with(
@@ -120,6 +123,7 @@ def test_search_params_dicts(self, client):
                 "limit": 10,
                 "external_id": "external_id",
                 "label": "label1,label2",
+                "location.saved_location_id": "saved_location_id",
             },
             verify=True,
         )
@@ -135,6 +139,7 @@ def test_create_params_underscores(self, client):
             location__unit="km",
             location__google_place_id="google_place_id",
             location__geoscope_paths=["geoscope_path1", "geoscope_path2"],
+            location__saved_location_id="saved_location_id",
             rank__type="type",
             rank__levels__phq={"min": 1.0, "max": 2.0},
             rank__levels__local={"min": 3.0, "max": 4.0},
@@ -158,6 +163,7 @@ def test_create_params_underscores(self, client):
                     "unit": "km",
                     "google_place_id": "google_place_id",
                     "geoscope_paths": ["geoscope_path1", "geoscope_path2"],
+                    "saved_location_id": "saved_location_id",
                 },
                 "rank": {
                     "type": "type",
@@ -192,6 +198,7 @@ def test_create_params_dicts(self, client):
                 "unit": "km",
                 "google_place_id": "google_place_id",
                 "geoscope_paths": ["geoscope_path1", "geoscope_path2"],
+                "saved_location_id": "saved_location_id",
             },
             rank={
                 "type": "type",
@@ -222,6 +229,7 @@ def test_create_params_dicts(self, client):
                     "unit": "km",
                     "google_place_id": "google_place_id",
                     "geoscope_paths": ["geoscope_path1", "geoscope_path2"],
+                    "saved_location_id": "saved_location_id",
                 },
                 "rank": {
                     "type": "type",
@@ -255,6 +263,7 @@ def test_update_params_underscores(self, client):
             location__unit="km",
             location__google_place_id="google_place_id",
             location__geoscope_paths=["geoscope_path1", "geoscope_path2"],
+            location__saved_location_id="saved_location_id",
             rank__type="type",
             rank__levels__phq={"min": 1.0, "max": 2.0},
             rank__levels__local={"min": 3.0, "max": 4.0},
@@ -278,6 +287,7 @@ def test_update_params_underscores(self, client):
                     "unit": "km",
                     "google_place_id": "google_place_id",
                     "geoscope_paths": ["geoscope_path1", "geoscope_path2"],
+                    "saved_location_id": "saved_location_id",
                 },
                 "rank": {
                     "type": "type",
@@ -310,6 +320,7 @@ def test_update_params_dicts(self, client):
                 "unit": "km",
                 "google_place_id": "google_place_id",
                 "geoscope_paths": ["geoscope_path1", "geoscope_path2"],
+                "saved_location_id": "saved_location_id",
             },
             rank={
                 "type": "type",
@@ -340,6 +351,7 @@ def test_update_params_dicts(self, client):
                     "unit": "km",
                     "google_place_id": "google_place_id",
                     "geoscope_paths": ["geoscope_path1", "geoscope_path2"],
+                    "saved_location_id": "saved_location_id",
                 },
                 "rank": {
                     "type": "type",