diff --git a/CHANGELOG.md b/CHANGELOG.md index d4a0c03eb1..951a4fce74 100644 --- a/CHANGELOG.md +++ b/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)) diff --git a/pybamm/parameters/parameter_values.py b/pybamm/parameters/parameter_values.py index dd596e746f..359033d7f8 100644 --- a/pybamm/parameters/parameter_values.py +++ b/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 diff --git a/pybamm/util.py b/pybamm/util.py index 885d2421b3..f58e04280f 100644 --- a/pybamm/util.py +++ b/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): """ diff --git a/tests/unit/test_util.py b/tests/unit/test_util.py index 646cc33af5..b4b7f6203e 100644 --- a/tests/unit/test_util.py +++ b/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