Move adstock and saturation method imports to mmm.__all__ (pymc-labs#908

) * Resolves pymc-labs#892: Move adstock and saturation method imports to mmm.__all__ * fix: patch uml GHA * Update .github/workflows/uml.yml * Update .github/workflows/uml.yml * Add permissions: write-all to .github/workflows/uml.yml --------- Co-authored-by: Will Dean <[email protected]>
PabloRoque · Aug 5, 2024 · 39f2336 · 39f2336
1 parent 41dd8a3
commit 39f2336
Show file tree

Hide file tree

Showing 5 changed files with 24 additions and 6 deletions.
diff --git a/.github/workflows/uml.yml b/.github/workflows/uml.yml
@@ -10,6 +10,8 @@ on:
     paths:
       - "pymc_marketing/**"
 
+permissions: write-all
+
 jobs:
   build:
     runs-on: ubuntu-latest
@@ -39,6 +41,9 @@ jobs:
           if git diff --staged --exit-code; then
             echo "No changes to commit"
           else
+            echo "Committing the changes"
             git commit -m "Update UML Diagrams"
-            git push
+            git push origin HEAD:${GITHUB_HEAD_REF}
           fi
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
diff --git a/environment.yml b/environment.yml
@@ -12,6 +12,8 @@ dependencies:
 - pandas
 - streamlit>=1.25.0
 - pip
+- pydantic
+- preliz
 # NOTE: Keep minimum pymc version in sync with ci.yml `OLDEST_PYMC_VERSION`
 - pymc>=5.12.0,<5.16.0
 - scikit-learn>=1.1.1
@@ -41,3 +43,4 @@ dependencies:
 - lifetimes==0.11.3
 - pytest==7.0.1
 - pytest-cov==3.0.0
+- pytest-mock
diff --git a/pymc_marketing/mmm/__init__.py b/pymc_marketing/mmm/__init__.py
@@ -20,6 +20,8 @@
     WeibullAdstock,
     WeibullCDFAdstock,
     WeibullPDFAdstock,
+    adstock_from_dict,
+    register_adstock_transformation,
 )
 from pymc_marketing.mmm.components.saturation import (
     HillSaturation,
@@ -30,6 +32,8 @@
     SaturationTransformation,
     TanhSaturation,
     TanhSaturationBaselined,
+    register_saturation_transformation,
+    saturation_from_dict,
 )
 from pymc_marketing.mmm.delayed_saturated_mmm import MMM, DelayedSaturatedMMM
 from pymc_marketing.mmm.fourier import MonthlyFourier, YearlyFourier
@@ -56,9 +60,13 @@
     "SaturationTransformation",
     "TanhSaturation",
     "TanhSaturationBaselined",
+    "saturation_from_dict",
+    "register_saturation_transformation",
     "WeibullAdstock",
     "WeibullCDFAdstock",
     "WeibullPDFAdstock",
+    "adstock_from_dict",
+    "register_adstock_transformation",
     "YearlyFourier",
     "base",
     "delayed_saturated_mmm",

diff --git a/tests/mmm/components/test_adstock.py b/tests/mmm/components/test_adstock.py
@@ -20,18 +20,20 @@
 import xarray as xr
 from pydantic import ValidationError
 
-from pymc_marketing.mmm.components.adstock import (
-    ADSTOCK_TRANSFORMATIONS,
+from pymc_marketing.mmm import (
     AdstockTransformation,
     DelayedAdstock,
     GeometricAdstock,
     WeibullAdstock,
     WeibullCDFAdstock,
     WeibullPDFAdstock,
-    _get_adstock_function,
     adstock_from_dict,
     register_adstock_transformation,
 )
+from pymc_marketing.mmm.components.adstock import (
+    ADSTOCK_TRANSFORMATIONS,
+    _get_adstock_function,
+)
 from pymc_marketing.mmm.transformers import ConvMode
 from pymc_marketing.prior import Prior
 

diff --git a/tests/mmm/components/test_saturation.py b/tests/mmm/components/test_saturation.py
@@ -20,17 +20,17 @@
 import xarray as xr
 from pydantic import ValidationError
 
-from pymc_marketing.mmm.components.saturation import (
+from pymc_marketing.mmm import (
     HillSaturation,
     InverseScaledLogisticSaturation,
     LogisticSaturation,
     MichaelisMentenSaturation,
     RootSaturation,
     TanhSaturation,
     TanhSaturationBaselined,
-    _get_saturation_function,
     saturation_from_dict,
 )
+from pymc_marketing.mmm.components.saturation import _get_saturation_function
 from pymc_marketing.prior import Prior