Deserialize arbitrary objects (#1296)

* support for deserialization of three classes

* add the deserialize logic

* correct the type hint

* add an example

* separate out the error

* catch error at deserialization

* relax the input type

* add test suite

* add test for arbitrary serialization via Prior

* use deserialize within from_json

* test for deserialize support within Prior

* use general deserialization in individual media transformations

* test both deserialize funcs

* test arb deserialization in adstock

* add similar deserialization check for saturation

* better naming of the tests

* support parsing of hsgp kwargs

* add to the module level docstring

* allow VariableFactory in parse_model_config

* add module to documentation

* add to the module documentation

* Reorder the top level documentation

* add more docstring about the functions

Author: williambdean

docs/source/api/index.md

@@ -15,4 +15,5 @@
   prior
   metrics
   mlflow
+  deserialize
 ```

pymc_marketing/deserialize.py

@@ -0,0 +1,244 @@
+#   Copyright 2024 The PyMC Labs Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+"""Deserialize into a PyMC-Marketing object.
+
+This is a two step process:
+
+1. Determine if the data is of the correct type.
+2. Deserialize the data into a python object for PyMC-Marketing.
+
+This is used to deserialize JSON data into PyMC-Marketing objects
+throughout the package.
+
+Examples
+--------
+Make use of the already registered PyMC-Marketing deserializers:
+
+.. code-block:: python
+
+    from pymc_marketing.deserialize import deserialize
+
+    prior_class_data = {
+        "dist": "Normal",
+        "kwargs": {"mu": 0, "sigma": 1}
+    }
+    prior = deserialize(prior_class_data)
+    # Prior("Normal", mu=0, sigma=1)
+
+Register custom class deserialization:
+
+.. code-block:: python
+
+    from pymc_marketing.deserialize import register_deserialization
+
+    class MyClass:
+        def __init__(self, value: int):
+            self.value = value
+
+        def to_dict(self) -> dict:
+            # Example of what the to_dict method might look like.
+            return {"value": self.value}
+
+    register_deserialization(
+        is_type=lambda data: data.keys() == {"value"} and isinstance(data["value"], int),
+        deserialize=lambda data: MyClass(value=data["value"]),
+    )
+
+Deserialize data into that custom class:
+
+.. code-block:: python
+
+    from pymc_marketing.deserialize import deserialize
+
+    data = {"value": 42}
+    obj = deserialize(data)
+    assert isinstance(obj, MyClass)
+
+
+"""
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+IsType = Callable[[Any], bool]
+Deserialize = Callable[[Any], Any]
+
+
+@dataclass
+class Deserializer:
+    """Object to store information required for deserialization.
+
+    All deserializers should be stored via the :func:`register_deserialization` function
+    instead of creating this object directly.
+
+    Attributes
+    ----------
+    is_type : IsType
+        Function to determine if the data is of the correct type.
+    deserialize : Deserialize
+        Function to deserialize the data.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        from typing import Any
+
+        class MyClass:
+            def __init__(self, value: int):
+                self.value = value
+
+        from pymc_marketing.deserialize import Deserializer
+
+        def is_type(data: Any) -> bool:
+            return data.keys() == {"value"} and isinstance(data["value"], int)
+
+        def deserialize(data: dict) -> MyClass:
+            return MyClass(value=data["value"])
+
+        deserialize_logic = Deserializer(is_type=is_type, deserialize=deserialize)
+
+    """
+
+    is_type: IsType
+    deserialize: Deserialize
+
+
+DESERIALIZERS: list[Deserializer] = []
+
+
+class DeserializableError(Exception):
+    """Error raised when data cannot be deserialized."""
+
+    def __init__(self, data: Any):
+        self.data = data
+        super().__init__(
+            f"Couldn't deserialize {data}. Use register_deserialization to add a deserialization mapping."
+        )
+
+
+def deserialize(data: Any) -> Any:
+    """Deserialize a dictionary into a Python object.
+
+    Use the :func:`register_deserialization` function to add custom deserializations.
+
+    Deserialization is a two step process due to the dynamic nature of the data:
+
+    1. Determine if the data is of the correct type.
+    2. Deserialize the data into a Python object.
+
+    Each registered deserialization is checked in order until one is found that can
+    deserialize the data. If no deserialization is found, a :class:`DeserializableError` is raised.
+
+    A :class:`DeserializableError` is raised when the data fails to be deserialized
+    by any of the registered deserializers.
+
+    Parameters
+    ----------
+    data : Any
+        The data to deserialize.
+
+    Returns
+    -------
+    Any
+        The deserialized object.
+
+    Raises
+    ------
+    DeserializableError
+        Raised when the data doesn't match any registered deserializations
+        or fails to be deserialized.
+
+    Examples
+    --------
+    Deserialize a :class:`pymc_marketing.prior.Prior` object:
+
+    .. code-block:: python
+
+        from pymc_marketing.deserialize import deserialize
+
+        data = {"dist": "Normal", "kwargs": {"mu": 0, "sigma": 1}}
+        prior = deserialize(data)
+        # Prior("Normal", mu=0, sigma=1)
+
+    """
+    for mapping in DESERIALIZERS:
+        try:
+            is_type = mapping.is_type(data)
+        except Exception:
+            is_type = False
+
+        if not is_type:
+            continue
+
+        try:
+            return mapping.deserialize(data)
+        except Exception as e:
+            raise DeserializableError(data) from e
+    else:
+        raise DeserializableError(data)
+
+
+def register_deserialization(is_type: IsType, deserialize: Deserialize) -> None:
+    """Register an arbitrary deserialization.
+
+    Use the :func:`deserialize` function to then deserialize data using all registered
+    deserialize functions.
+
+    Classes from PyMC-Marketing have their deserialization mappings registered
+    automatically. However, custom classes will need to be registered manually
+    using this function before they can be deserialized.
+
+    Parameters
+    ----------
+    is_type : Callable[[Any], bool]
+        Function to determine if the data is of the correct type.
+    deserialize : Callable[[dict], Any]
+        Function to deserialize the data of that type.
+
+    Examples
+    --------
+    Register a custom class deserialization:
+
+    .. code-block:: python
+
+        from pymc_marketing.deserialize import register_deserialization
+
+        class MyClass:
+            def __init__(self, value: int):
+                self.value = value
+
+            def to_dict(self) -> dict:
+                # Example of what the to_dict method might look like.
+                return {"value": self.value}
+
+        register_deserialization(
+            is_type=lambda data: data.keys() == {"value"} and isinstance(data["value"], int),
+            deserialize=lambda data: MyClass(value=data["value"]),
+        )
+
+    Use that custom class deserialization:
+
+    .. code-block:: python
+
+        from pymc_marketing.deserialize import deserialize
+
+        data = {"value": 42}
+        obj = deserialize(data)
+        assert isinstance(obj, MyClass)
+
+    """
+    mapping = Deserializer(is_type=is_type, deserialize=deserialize)
+    DESERIALIZERS.append(mapping)

pymc_marketing/hsgp_kwargs.py

@@ -18,6 +18,8 @@
 import pymc as pm
 from pydantic import BaseModel, Field, InstanceOf
 
+from pymc_marketing.deserialize import register_deserialization
+
 
 class HSGPKwargs(BaseModel):
     """HSGP keyword arguments for the time-varying prior.
@@ -80,3 +82,20 @@ class HSGPKwargs(BaseModel):
     cov_func: InstanceOf[pm.gp.cov.Covariance] | str | None = Field(
         None, description="Gaussian process Covariance function"
     )
+
+
+def _is_hsgp_kwargs(data) -> bool:
+    return isinstance(data, dict) and data.keys() == {
+        "m",
+        "L",
+        "eta_lam",
+        "ls_mu",
+        "ls_sigma",
+        "cov_func",
+    }
+
+
+register_deserialization(
+    is_type=_is_hsgp_kwargs,
+    deserialize=lambda data: HSGPKwargs.model_validate(data),
+)

pymc_marketing/mmm/components/adstock.py

@@ -56,10 +56,10 @@ def function(self, x, alpha):
 import xarray as xr
 from pydantic import Field, validate_call
 
+from pymc_marketing.deserialize import deserialize, register_deserialization
 from pymc_marketing.mmm.components.base import (
     SupportedPrior,
     Transformation,
-    _deserialize,
 )
 from pymc_marketing.mmm.transformers import (
     ConvMode,
@@ -345,6 +345,16 @@ def adstock_from_dict(data: dict) -> AdstockTransformation:
     cls = ADSTOCK_TRANSFORMATIONS[lookup_name]
 
     if "priors" in data:
-        data["priors"] = {k: _deserialize(v) for k, v in data["priors"].items()}
+        data["priors"] = {k: deserialize(v) for k, v in data["priors"].items()}
 
     return cls(**data)
+
+
+def _is_adstock(data):
+    return "lookup_name" in data and data["lookup_name"] in ADSTOCK_TRANSFORMATIONS
+
+
+register_deserialization(
+    is_type=_is_adstock,
+    deserialize=adstock_from_dict,
+)

pymc_marketing/mmm/components/base.py

@@ -592,10 +592,3 @@ def _serialize_value(value: Any) -> Any:
         return value.tolist()
 
     return value
-
-
-def _deserialize(value):
-    try:
-        return Prior.from_json(value)
-    except Exception:
-        return value

pymc_marketing/mmm/components/saturation.py

@@ -76,9 +76,9 @@ def function(self, x, b):
 import xarray as xr
 from pydantic import Field, InstanceOf, validate_call
 
+from pymc_marketing.deserialize import deserialize, register_deserialization
 from pymc_marketing.mmm.components.base import (
     Transformation,
-    _deserialize,
 )
 from pymc_marketing.mmm.transformers import (
     hill_function,
@@ -483,6 +483,13 @@ def saturation_from_dict(data: dict) -> SaturationTransformation:
 
     if "priors" in data:
         data["priors"] = {
-            key: _deserialize(value) for key, value in data["priors"].items()
+            key: deserialize(value) for key, value in data["priors"].items()
         }
     return cls(**data)
+
+
+def _is_saturation(data):
+    return "lookup_name" in data and data["lookup_name"] in SATURATION_TRANSFORMATIONS
+
+
+register_deserialization(_is_saturation, saturation_from_dict)

pymc_marketing/model_config.py

@@ -16,15 +16,16 @@
 import warnings
 from typing import Any
 
+from pymc_marketing.deserialize import deserialize
 from pymc_marketing.hsgp_kwargs import HSGPKwargs
-from pymc_marketing.prior import Prior
+from pymc_marketing.prior import Prior, VariableFactory
 
 
 class ModelConfigError(Exception):
     """Exception raised for errors in model configuration."""
 
 
-ModelConfig = dict[str, HSGPKwargs | Prior | Any]
+ModelConfig = dict[str, VariableFactory | HSGPKwargs | Prior | Any]
 
 
 def parse_model_config(
@@ -121,11 +122,11 @@ def handle_prior_config(name, prior_config):
         if name in non_distributions or name in hsgp_kwargs_fields:
             return prior_config
 
-        if isinstance(prior_config, Prior):
+        if isinstance(prior_config, Prior) or isinstance(prior_config, VariableFactory):
             return prior_config
 
         try:
-            dist = Prior.from_json(prior_config)
+            dist = deserialize(prior_config)
         except Exception as e:
             parse_errors.append(f"Parameter {name}: {e}")
         else: