Merge pull request #875 from pybamm-team/issue-869-search

Issue 869 search

CHANGELOG.md

@@ -1,7 +1,8 @@
 # [Unreleased](https://github.com/pybamm-team/PyBaMM)
 
-## Features 
+## Features
 
+-   Added a search feature to `FuzzyDict` ([#875](https://github.com/pybamm-team/PyBaMM/pull/875))
 -   Add ambient temperature as a function of time ([#872](https://github.com/pybamm-team/PyBaMM/pull/872))
 -   Added `CasadiAlgebraicSolver` for solving algebraic systems with casadi ([#868](https://github.com/pybamm-team/PyBaMM/pull/868))
 

pybamm/parameters/parameter_values.py

@@ -103,6 +103,14 @@ def items(self):
         "Get the items of the dictionary"
         return self._dict_items.items()
 
+    def search(self, key, print_values=True):
+        """
+        Search dictionary for keys containing 'key'.
+
+        See :meth:`pybamm.FuzzyDict.search()`.
+        """
+        return self._dict_items.search(key, print_values)
+
     def update_from_chemistry(self, chemistry):
         """
         Load standard set of components from a 'chemistry' dictionary

pybamm/util.py

@@ -91,6 +91,37 @@ def __getitem__(self, key):
             best_matches = self.get_best_matches(key)
             raise KeyError(f"'{key}' not found. Best matches are {best_matches}")
 
+    def search(self, key, print_values=False):
+        """
+        Search dictionary for keys containing 'key'. If print_values is True, then
+        both the keys and values will be printed. Otherwise just the values will
+        be printed. If no results are found, the best matches are printed.
+        """
+        key = key.lower()
+
+        # Sort the keys so results are stored in alphabetical order
+        keys = list(self.keys())
+        keys.sort()
+        results = {}
+
+        # Check if any of the dict keys contain the key we are searching for
+        for k in keys:
+            if key in k.lower():
+                results[k] = self[k]
+
+        if results == {}:
+            # If no results, return best matches
+            best_matches = self.get_best_matches(key)
+            print(
+                f"No results for search using '{key}'. Best matches are {best_matches}"
+            )
+        elif print_values:
+            # Else print results, including dict items
+            print("\n".join("{}\t{}".format(k, v) for k, v in results.items()))
+        else:
+            # Just print keys
+            print("\n".join("{}".format(k) for k in results.keys()))
+
 
 class Timer(object):
     """

tests/unit/test_util.py

@@ -6,6 +6,8 @@
 import pybamm
 import tempfile
 import unittest
+from unittest.mock import patch
+from io import StringIO
 
 
 class TestUtil(unittest.TestCase):
@@ -97,6 +99,34 @@ def test_get_parameters_filepath(self):
         self.assertTrue(pybamm.get_parameters_filepath(tempfile_obj.name) == path)
 
 
+class TestSearch(unittest.TestCase):
+    def test_url_gets_to_stdout(self):
+        model = pybamm.BaseModel()
+        model.variables = {"Electrolyte concentration": 1, "Electrode potential": 0}
+
+        param = pybamm.ParameterValues({"a": 10, "b": 2})
+
+        # Test variables search (default returns key)
+        with patch("sys.stdout", new=StringIO()) as fake_out:
+            model.variables.search("Electrode")
+            self.assertEqual(fake_out.getvalue(), "Electrode potential\n")
+
+        # Test bad var search (returns best matches)
+        with patch("sys.stdout", new=StringIO()) as fake_out:
+            model.variables.search("bad var")
+            out = (
+                "No results for search using 'bad var'. "
+                "Best matches are ['Electrolyte concentration', "
+                "'Electrode potential']\n"
+            )
+            self.assertEqual(fake_out.getvalue(), out)
+
+        # Test param search (default returns key, value)
+        with patch("sys.stdout", new=StringIO()) as fake_out:
+            param.search("a")
+            self.assertEqual(fake_out.getvalue(), "a\t10\n")
+
+
 if __name__ == "__main__":
     print("Add -v for more debug output")
     import sys