Merge pull request #55 from zytedata/feature/location-params

Add LocationParam

.python-version

@@ -0,0 +1 @@
+3.9

tests/test_params_location_param.py

@@ -0,0 +1,52 @@
+import pytest
+from pydantic import ValidationError
+
+from zyte_spider_templates.params import LocationParam
+
+
+def test_valid_location_param():
+    valid_address_dict = {
+        "streetAddress": "123 Main St",
+        "addressCountry": "US",
+        "addressRegion": "CA",
+        "postalCode": "12345",
+    }
+    location_param = LocationParam(location=valid_address_dict)  # type: ignore[arg-type]
+    assert location_param.location is not None
+    assert location_param.location.streetAddress == "123 Main St"
+    assert location_param.location.addressCountry == "US"
+    assert location_param.location.addressRegion == "CA"
+    assert location_param.location.postalCode == "12345"
+
+
+def test_valid_location_param_from_json():
+    valid_address_json = '{"streetAddress": "456 Elm St", "addressCountry": "US", "addressRegion": "NY", "postalCode": "54321"}'
+    location_param = LocationParam(location=valid_address_json)  # type: ignore[arg-type]
+    assert location_param.location is not None
+    assert location_param.location.streetAddress == "456 Elm St"
+    assert location_param.location.addressCountry == "US"
+    assert location_param.location.addressRegion == "NY"
+    assert location_param.location.postalCode == "54321"
+
+
+def test_none_location_param():
+    location_param = LocationParam(location=None)
+    assert location_param.location is None
+
+
+def test_invalid_json_location_param():
+    invalid_address_json = '{"streetAddress": "789 Pine St", "addressCountry": "AnotheraddressCountry", "addressRegion": "FL", "postalCode": "67890"'
+    with pytest.raises(ValueError, match=r".* is not a valid JSON object"):
+        LocationParam(location=invalid_address_json)  # type: ignore[arg-type]
+
+
+def test_invalid_type_location_param():
+    invalid_type_value = 12345  # Invalid type, should raise ValueError
+    with pytest.raises(ValueError, match=r".* type .* is not a supported type"):
+        LocationParam(location=invalid_type_value)  # type: ignore[arg-type]
+
+
+def test_invalid_validation_location_param():
+    invalid_address_json = '{"nonExpectedInputField": "67890"}'
+    with pytest.raises(ValidationError, match=r"Extra inputs are not permitted .*"):
+        LocationParam(location=invalid_address_json)  # type: ignore[arg-type]

zyte_spider_templates/params.py

@@ -1,7 +1,8 @@
+import json
 from enum import Enum
-from typing import Optional
+from typing import Dict, Optional, Union
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field, field_validator
 
 from zyte_spider_templates._geolocations import (
     GEOLOCATION_OPTIONS_WITH_CODE,
@@ -107,3 +108,68 @@ class UrlParam(BaseModel):
             "exclusiveRequired": True,
         },
     )
+
+
+class PostalAddress(BaseModel):
+    """
+    Represents a postal address with various optional components such as
+    street address, postal code, region, and country.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    streetAddress: Optional[str] = Field(
+        title="Street Address",
+        description="The street address",
+        default=None,
+    )
+
+    postalCode: Optional[str] = Field(
+        title="Postal Code",
+        description="The postal code",
+        default=None,
+    )
+
+    addressRegion: Optional[str] = Field(
+        title="Address Region",
+        description="The region in which the address is. This value is specific to the website",
+        default=None,
+    )
+
+    addressCountry: Optional[str] = Field(
+        title="Adress Country",
+        description="The country code in ISO 3166-1 alpha-2",
+        default=None,
+    )
+
+
+class LocationParam(BaseModel):
+    """
+    Represents a parameter containing a postal address to be set as the user location on websites.
+    """
+
+    location: Optional[PostalAddress] = Field(
+        title="Location",
+        description="Postal address to be set as the user location on websites",
+        default=None,
+    )
+
+    @field_validator("location", mode="before")
+    @classmethod
+    def validate_location(
+        cls, value: Optional[Union[PostalAddress, str, Dict]]
+    ) -> Optional[PostalAddress]:
+        """Validate location field and cast it into PostalAddress if needed"""
+        if value is None or isinstance(value, PostalAddress):
+            return value
+
+        if isinstance(value, str):
+            try:
+                return PostalAddress(**json.loads(value))
+            except json.decoder.JSONDecodeError as err:
+                raise ValueError(f"{value!r} is not a valid JSON object") from err
+
+        elif isinstance(value, dict):
+            return PostalAddress(**value)
+
+        raise ValueError(f"{value!r} type {type(value)} is not a supported type")