Unit tests for partial_format

composer/utils/string_helpers.py

@@ -24,7 +24,9 @@ def partial_format(s, *args, **kwargs):
         try:
             result = s.format(*args, **kwargs)
             done = True
-        except KeyError as e:
+        except IndexError as e:  # Missing positional arg
+            args += ('{}',)
+        except KeyError as e:  # Missing keyword arg
             key = e.args[0]
             kwargs[key] = '{' + key + '}'
 

tests/utils/test_string_helpers.py

@@ -0,0 +1,21 @@
+# Copyright 2022 MosaicML Composer authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Unit tests for string helpers."""
+
+from composer.utils.string_helpers import partial_format
+
+
+def test_partial_format():
+    # Keyword args
+    assert partial_format('{foo} {bar}', foo='Hello') == 'Hello {bar}'
+    assert partial_format('{foo} {bar}', foo='Hello', bar='World') == 'Hello World'
+
+    # Positional args
+    assert partial_format('{} {}', 'Hello') == 'Hello {}'
+    assert partial_format('{} {}', 'Hello', 'World') == 'Hello World'
+
+    # Positional and keyword args
+    assert partial_format('{foo} {}', 'World') == '{foo} World'
+    assert partial_format('{foo} {}', foo='Hello') == 'Hello {}'
+    assert partial_format('{foo} {}', 'World', foo='Hello') == 'Hello World'