✨ ⬆️ Make kedro-mlflow compatible with pydantic v2 (#476)

CHANGELOG.md

@@ -5,6 +5,7 @@
 ### Added
 
 - :sparkles: Add support for python 3.11 ([#450, rxm7706](https://github.com/Galileo-Galilei/kedro-mlflow/pull/450))
+- :sparkles: :arrow_up: Add support for pydantic v2 ([#476](https://github.com/Galileo-Galilei/kedro-mlflow/pull/476))
 
 ### Changed
 

kedro_mlflow/framework/hooks/mlflow_hook.py

@@ -16,6 +16,7 @@
 from mlflow.models import infer_signature
 from mlflow.tracking import MlflowClient
 from mlflow.utils.validation import MAX_PARAM_VAL_LENGTH
+from pydantic import __version__ as pydantic_version
 
 from kedro_mlflow.config.kedro_mlflow_config import KedroMlflowConfig
 from kedro_mlflow.framework.hooks.utils import (
@@ -74,7 +75,11 @@ def after_context_created(
             # but we got an empty dict
             conf_mlflow_yml = {}
 
-        mlflow_config = KedroMlflowConfig.parse_obj({**conf_mlflow_yml})
+        mlflow_config = (
+            KedroMlflowConfig.model_validate({**conf_mlflow_yml})
+            if pydantic_version > "2.0.0"
+            else KedroMlflowConfig.parse_obj({**conf_mlflow_yml})
+        )
 
         self._already_active_mlflow = False
         if mlflow.active_run():

requirements.txt

@@ -1,4 +1,4 @@
 kedro>=0.18.1, <0.19.0
 kedro_datasets
 mlflow>=1.0.0, <3.0.0
-pydantic>=1.0.0, <2.0.0
+pydantic>=1.0.0, <3.0.0

tests/config/test_get_mlflow_config.py

@@ -360,7 +360,7 @@ def request_headers(self):
             credentials=None,
             request_header_provider=dict(
                 type="custom_rhp.CustomRequestHeaderProviderInitKwargs",
-                init_kwargs=dict(a=1),
+                init_kwargs=dict(a="a"),
             ),
         ),
         tracking=dict(
@@ -389,7 +389,7 @@ def request_headers(self):
             mtrr._request_header_provider_registry._registry[-1].__class__.__name__
             == "CustomRequestHeaderProviderInitKwargs"
         )
-    assert mtrr._request_header_provider_registry._registry[-1].a == "1"
+    assert mtrr._request_header_provider_registry._registry[-1].a == "a"
     assert not hasattr(mtrr._request_header_provider_registry._registry[-1], "context")
 
 
@@ -423,7 +423,7 @@ def request_headers(self):
             request_header_provider=dict(
                 type="custom_rhp.CustomRequestHeaderProviderInitKwargsKedroContext",
                 pass_context=True,
-                init_kwargs=dict(b=2),
+                init_kwargs=dict(b="b"),
             ),
         ),
         tracking=dict(
@@ -452,7 +452,7 @@ def request_headers(self):
         mtrr._request_header_provider_registry._registry[-1].__class__.__name__
         == "CustomRequestHeaderProviderInitKwargsKedroContext"
     )
-    assert mtrr._request_header_provider_registry._registry[-1].b == "2"
+    assert mtrr._request_header_provider_registry._registry[-1].b == "b"
     assert hasattr(mtrr._request_header_provider_registry._registry[-1], "context")
     assert isinstance(
         mtrr._request_header_provider_registry._registry[-1].context, KedroContext