diff --git a/src/pybamm/experiment/step/base_step.py b/src/pybamm/experiment/step/base_step.py index a7dfa9c9ba..8dd829935e 100644 --- a/src/pybamm/experiment/step/base_step.py +++ b/src/pybamm/experiment/step/base_step.py @@ -1,6 +1,4 @@ -# -# Private classes and functions for experiment steps -# +from __future__ import annotations import pybamm import numpy as np from datetime import datetime @@ -492,6 +490,10 @@ def set_up(self, new_model, new_parameter_values): } +def get_unit_from(a_string: str) -> str: + return a_string.lstrip("0123456789.-eE ") + + def _convert_time_to_seconds(time_and_units): """Convert a time in seconds, minutes or hours to a time in seconds""" if time_and_units is None: @@ -501,11 +503,10 @@ def _convert_time_to_seconds(time_and_units): if isinstance(time_and_units, numbers.Number): if time_and_units <= 0: raise ValueError("time must be positive") - else: - return time_and_units + return time_and_units # Split number and units - units = time_and_units.lstrip("0123456789.- ") + units = get_unit_from(time_and_units) time = time_and_units[: -len(units)] if units in ["second", "seconds", "s", "sec"]: time_in_seconds = float(time) @@ -528,7 +529,7 @@ def _convert_temperature_to_kelvin(temperature_and_units): return temperature_and_units # Split number and units - units = temperature_and_units.lstrip("0123456789.- ") + units = get_unit_from(temperature_and_units) temperature = temperature_and_units[: -len(units)] if units in ["K"]: temperature_in_kelvin = float(temperature) @@ -547,7 +548,7 @@ def _convert_electric(value_string): value = 1 / float(value_string[2:]) else: # All other cases e.g. 4 A, 2.5 V, 1.5 Ohm - unit = value_string.lstrip("0123456789.- ") + unit = get_unit_from(value_string) value = float(value_string[: -len(unit)]) # Catch milli- prefix if unit.startswith("m"): diff --git a/tests/unit/test_experiments/test_base_step.py b/tests/unit/test_experiments/test_base_step.py new file mode 100644 index 0000000000..0250c0622f --- /dev/null +++ b/tests/unit/test_experiments/test_base_step.py @@ -0,0 +1,17 @@ +import pytest +import pybamm + + +@pytest.mark.parametrize( + "test_string, unit_string", + [ + ("123e-1 W", "W"), + ("123K", "K"), + ("1A", "A"), + ("2.0 mV", "mV"), + ("0.5 Ohm", "Ohm"), + ("1e0hours", "hours"), + ], +) +def test_read_units(test_string, unit_string): + assert unit_string == pybamm.experiment.step.base_step.get_unit_from(test_string)