From e148006fc46f75edbddeec07ee943434793daa62 Mon Sep 17 00:00:00 2001
From: Abhishek Chaudhari
 <91185083+AbhishekChaudharii@users.noreply.github.com>
Date: Wed, 25 Sep 2024 22:01:34 +0530
Subject: [PATCH] Added __iter__ and __contains__ method to
 pybamm.ParameterValues (#4465)

* Added __contains__ and __iter__ in pybamm.ParameterValues

* After running pre_commit

* added tests to verify __iter__ and __contains__ method work-2

* Update tests/unit/test_parameters/test_parameter_values.py

Deleted print statement which was intended for debugging

Co-authored-by: Agriya Khetarpal <74401230+agriyakhetarpal@users.noreply.github.com>

* Update tests/unit/test_parameters/test_parameter_values.py

---------

Co-authored-by: Agriya Khetarpal <74401230+agriyakhetarpal@users.noreply.github.com>
Co-authored-by: Arjun Verma <arjunverma.oc@gmail.com>
---
 src/pybamm/parameters/parameter_values.py     |  6 ++++++
 .../test_parameters/test_parameter_values.py  | 21 +++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/src/pybamm/parameters/parameter_values.py b/src/pybamm/parameters/parameter_values.py
index 43c1ea17ce..d5d5878cf2 100644
--- a/src/pybamm/parameters/parameter_values.py
+++ b/src/pybamm/parameters/parameter_values.py
@@ -930,3 +930,9 @@ def print_evaluated_parameters(self, evaluated_parameters, output_file):
                     file.write((s + " : {:10.4g}\n").format(name, value))
                 else:
                     file.write((s + " : {:10.3E}\n").format(name, value))
+
+    def __contains__(self, key):
+        return key in self._dict_items
+
+    def __iter__(self):
+        return iter(self._dict_items)
diff --git a/tests/unit/test_parameters/test_parameter_values.py b/tests/unit/test_parameters/test_parameter_values.py
index 4086abea6d..543d3a6b4c 100644
--- a/tests/unit/test_parameters/test_parameter_values.py
+++ b/tests/unit/test_parameters/test_parameter_values.py
@@ -17,6 +17,7 @@
 )
 from pybamm.expression_tree.exceptions import OptionError
 import casadi
+from pybamm.parameters.parameter_values import ParameterValues
 
 
 class TestParameterValues:
@@ -1029,3 +1030,23 @@ def test_exchange_current_density_plating(self):
             match="referring to the reaction at the surface of a lithium metal electrode",
         ):
             parameter_values.evaluate(param)
+
+    def test_contains_method(self):
+        """Test for __contains__ method to check the functionality of 'in' keyword"""
+        parameter_values = ParameterValues(
+            {"Negative particle radius [m]": 1e-6, "Positive particle radius [m]": 2e-6}
+        )
+        assert (
+            "Negative particle radius [m]" in parameter_values
+        ), "Key should be found in parameter_values"
+        assert (
+            "Invalid key" not in parameter_values
+        ), "Non-existent key should not be found"
+
+    def test_iter_method(self):
+        """Test for __iter__ method to check if we can iterate over keys"""
+        parameter_values = ParameterValues(
+            values={"Negative particle radius [m]": 1e-6}
+        )
+        pv = [i for i in parameter_values]
+        assert len(pv) == 5, "Should have 5 keys"